ob-metaflow-stubs 6.0.3.134__py2.py3-none-any.whl → 6.0.3.136__py2.py3-none-any.whl

Sign up to get free protection for your applications and to get access to all the features.
Files changed (192) hide show
  1. metaflow-stubs/__init__.pyi +231 -231
  2. metaflow-stubs/cards.pyi +2 -2
  3. metaflow-stubs/cli.pyi +2 -2
  4. metaflow-stubs/cli_components/__init__.pyi +2 -2
  5. metaflow-stubs/cli_components/utils.pyi +2 -2
  6. metaflow-stubs/client/__init__.pyi +2 -2
  7. metaflow-stubs/client/core.pyi +6 -6
  8. metaflow-stubs/client/filecache.pyi +3 -3
  9. metaflow-stubs/events.pyi +2 -2
  10. metaflow-stubs/exception.pyi +2 -2
  11. metaflow-stubs/flowspec.pyi +5 -5
  12. metaflow-stubs/generated_for.txt +1 -1
  13. metaflow-stubs/includefile.pyi +4 -4
  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 +104 -104
  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 +2 -2
  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 +3 -3
  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 +3 -3
  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 +3 -3
  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 +2 -2
  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 +4 -4
  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 +4 -4
  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 +4 -4
  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 +3 -3
  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 +3 -3
  72. metaflow-stubs/plugins/__init__.pyi +9 -9
  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 +3 -3
  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 +4 -4
  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 +4 -4
  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 +4 -4
  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 +3 -3
  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 +2 -2
  115. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +2 -2
  116. metaflow-stubs/plugins/cards/card_modules/basic.pyi +2 -2
  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 +4 -4
  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 +3 -3
  146. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +2 -2
  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/kubernetes/spot_monitor_sidecar.pyi +2 -2
  151. metaflow-stubs/plugins/parallel_decorator.pyi +2 -2
  152. metaflow-stubs/plugins/perimeters.pyi +2 -2
  153. metaflow-stubs/plugins/project_decorator.pyi +2 -2
  154. metaflow-stubs/plugins/pypi/__init__.pyi +3 -3
  155. metaflow-stubs/plugins/pypi/conda_decorator.pyi +2 -2
  156. metaflow-stubs/plugins/pypi/conda_environment.pyi +5 -5
  157. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +2 -2
  158. metaflow-stubs/plugins/pypi/pypi_environment.pyi +2 -2
  159. metaflow-stubs/plugins/pypi/utils.pyi +2 -2
  160. metaflow-stubs/plugins/resources_decorator.pyi +2 -2
  161. metaflow-stubs/plugins/retry_decorator.pyi +2 -2
  162. metaflow-stubs/plugins/secrets/__init__.pyi +2 -2
  163. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +3 -3
  164. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +2 -2
  165. metaflow-stubs/plugins/snowflake/__init__.pyi +2 -2
  166. metaflow-stubs/plugins/storage_executor.pyi +2 -2
  167. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +2 -2
  168. metaflow-stubs/plugins/timeout_decorator.pyi +3 -3
  169. metaflow-stubs/profilers/__init__.pyi +2 -2
  170. metaflow-stubs/pylint_wrapper.pyi +2 -2
  171. metaflow-stubs/runner/__init__.pyi +2 -2
  172. metaflow-stubs/runner/deployer.pyi +29 -29
  173. metaflow-stubs/runner/deployer_impl.pyi +2 -2
  174. metaflow-stubs/runner/metaflow_runner.pyi +5 -5
  175. metaflow-stubs/runner/nbdeploy.pyi +2 -2
  176. metaflow-stubs/runner/nbrun.pyi +2 -2
  177. metaflow-stubs/runner/subprocess_manager.pyi +2 -2
  178. metaflow-stubs/runner/utils.pyi +4 -4
  179. metaflow-stubs/system/__init__.pyi +2 -2
  180. metaflow-stubs/system/system_logger.pyi +2 -2
  181. metaflow-stubs/system/system_monitor.pyi +2 -2
  182. metaflow-stubs/tagging_util.pyi +2 -2
  183. metaflow-stubs/tuple_util.pyi +2 -2
  184. metaflow-stubs/user_configs/__init__.pyi +2 -2
  185. metaflow-stubs/user_configs/config_decorators.pyi +4 -4
  186. metaflow-stubs/user_configs/config_options.pyi +3 -3
  187. metaflow-stubs/user_configs/config_parameters.pyi +5 -5
  188. {ob_metaflow_stubs-6.0.3.134.dist-info → ob_metaflow_stubs-6.0.3.136.dist-info}/METADATA +1 -1
  189. ob_metaflow_stubs-6.0.3.136.dist-info/RECORD +192 -0
  190. ob_metaflow_stubs-6.0.3.134.dist-info/RECORD +0 -192
  191. {ob_metaflow_stubs-6.0.3.134.dist-info → ob_metaflow_stubs-6.0.3.136.dist-info}/WHEEL +0 -0
  192. {ob_metaflow_stubs-6.0.3.134.dist-info → ob_metaflow_stubs-6.0.3.136.dist-info}/top_level.txt +0 -0
@@ -1,15 +1,15 @@
1
1
  ######################################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
- # MF version: 2.13.7.1+obcheckpoint(0.1.6);ob(v1) #
4
- # Generated on 2025-01-23T20:10:29.134629 #
3
+ # MF version: 2.13.9.1+obcheckpoint(0.1.6);ob(v1) #
4
+ # Generated on 2025-01-31T18:07:50.935996 #
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
 
@@ -147,112 +147,120 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
147
147
  ...
148
148
 
149
149
  @typing.overload
150
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
150
+ 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]]]:
151
151
  """
152
- Decorator prototype for all step decorators. This function gets specialized
153
- and imported for all decorators types by _import_plugin_decorators().
152
+ Enables loading / saving of models within a step.
154
153
  """
155
154
  ...
156
155
 
157
156
  @typing.overload
158
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
157
+ def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
159
158
  ...
160
159
 
161
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
160
+ @typing.overload
161
+ def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
162
+ ...
163
+
164
+ 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):
162
165
  """
163
- Decorator prototype for all step decorators. This function gets specialized
164
- and imported for all decorators types by _import_plugin_decorators().
166
+ Enables loading / saving of models within a step.
165
167
  """
166
168
  ...
167
169
 
168
170
  @typing.overload
169
- 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]]]:
171
+ def resources(*, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
170
172
  """
171
- Enables checkpointing for a step.
173
+ Specifies the resources needed when executing this step.
174
+
175
+ Use `@resources` to specify the resource requirements
176
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
177
+
178
+ You can choose the compute layer on the command line by executing e.g.
179
+ ```
180
+ python myflow.py run --with batch
181
+ ```
182
+ or
183
+ ```
184
+ python myflow.py run --with kubernetes
185
+ ```
186
+ which executes the flow on the desired system using the
187
+ requirements specified in `@resources`.
172
188
  """
173
189
  ...
174
190
 
175
191
  @typing.overload
176
- def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
192
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
177
193
  ...
178
194
 
179
195
  @typing.overload
180
- def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
181
- ...
182
-
183
- 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):
184
- """
185
- Enables checkpointing for a step.
186
- """
196
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
187
197
  ...
188
198
 
189
- @typing.overload
190
- 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]]]:
199
+ def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None):
191
200
  """
192
- Specifies a timeout for your step.
193
-
194
- This decorator is useful if this step may hang indefinitely.
201
+ Specifies the resources needed when executing this step.
195
202
 
196
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
197
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
198
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
203
+ Use `@resources` to specify the resource requirements
204
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
199
205
 
200
- Note that all the values specified in parameters are added together so if you specify
201
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
206
+ You can choose the compute layer on the command line by executing e.g.
207
+ ```
208
+ python myflow.py run --with batch
209
+ ```
210
+ or
211
+ ```
212
+ python myflow.py run --with kubernetes
213
+ ```
214
+ which executes the flow on the desired system using the
215
+ requirements specified in `@resources`.
202
216
  """
203
217
  ...
204
218
 
205
219
  @typing.overload
206
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
220
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
221
+ """
222
+ Internal decorator to support Fast bakery
223
+ """
207
224
  ...
208
225
 
209
226
  @typing.overload
210
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
227
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
211
228
  ...
212
229
 
213
- 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):
230
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
214
231
  """
215
- Specifies a timeout for your step.
216
-
217
- This decorator is useful if this step may hang indefinitely.
218
-
219
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
220
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
221
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
222
-
223
- Note that all the values specified in parameters are added together so if you specify
224
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
232
+ Internal decorator to support Fast bakery
225
233
  """
226
234
  ...
227
235
 
228
236
  @typing.overload
229
- 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]]]:
237
+ 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]]]:
230
238
  """
231
- Creates a human-readable report, a Metaflow Card, after this step completes.
239
+ Specifies the Conda environment for the step.
232
240
 
233
- Note that you may add multiple `@card` decorators in a step with different parameters.
241
+ Information in this decorator will augment any
242
+ attributes set in the `@conda_base` flow-level decorator. Hence,
243
+ you can use `@conda_base` to set packages required by all
244
+ steps and use `@conda` to specify step-specific overrides.
234
245
  """
235
246
  ...
236
247
 
237
248
  @typing.overload
238
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
249
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
239
250
  ...
240
251
 
241
252
  @typing.overload
242
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
253
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
243
254
  ...
244
255
 
245
- 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):
256
+ 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):
246
257
  """
247
- Creates a human-readable report, a Metaflow Card, after this step completes.
258
+ Specifies the Conda environment for the step.
248
259
 
249
- Note that you may add multiple `@card` decorators in a step with different parameters.
250
- """
251
- ...
252
-
253
- 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]]]:
254
- """
255
- Specifies that this step should execute on DGX cloud.
260
+ Information in this decorator will augment any
261
+ attributes set in the `@conda_base` flow-level decorator. Hence,
262
+ you can use `@conda_base` to set packages required by all
263
+ steps and use `@conda` to specify step-specific overrides.
256
264
  """
257
265
  ...
258
266
 
@@ -287,73 +295,90 @@ def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
287
295
  """
288
296
  ...
289
297
 
298
+ 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]]]:
299
+ """
300
+ Specifies that this step should execute on DGX cloud.
301
+ """
302
+ ...
303
+
290
304
  @typing.overload
291
- 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]]]:
305
+ 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]]]:
292
306
  """
293
- Specifies the PyPI packages for the step.
294
-
295
- Information in this decorator will augment any
296
- attributes set in the `@pyi_base` flow-level decorator. Hence,
297
- you can use `@pypi_base` to set packages required by all
298
- steps and use `@pypi` to specify step-specific overrides.
307
+ Enables checkpointing for a step.
299
308
  """
300
309
  ...
301
310
 
302
311
  @typing.overload
303
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
312
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
304
313
  ...
305
314
 
306
315
  @typing.overload
307
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
316
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
308
317
  ...
309
318
 
310
- 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):
319
+ 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):
311
320
  """
312
- Specifies the PyPI packages for the step.
313
-
314
- Information in this decorator will augment any
315
- attributes set in the `@pyi_base` flow-level decorator. Hence,
316
- you can use `@pypi_base` to set packages required by all
317
- steps and use `@pypi` to specify step-specific overrides.
321
+ Enables checkpointing for a step.
318
322
  """
319
323
  ...
320
324
 
321
325
  @typing.overload
322
- 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]]]:
326
+ 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]]]:
323
327
  """
324
- Specifies the number of times the task corresponding
325
- to a step needs to be retried.
328
+ Specifies environment variables to be set prior to the execution of a step.
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.
343
+ """
344
+ ...
345
+
346
+ @typing.overload
347
+ 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]]]:
348
+ """
349
+ Specifies a timeout for your step.
326
350
 
327
- This decorator is useful for handling transient errors, such as networking issues.
328
- If your task contains operations that can't be retried safely, e.g. database updates,
329
- it is advisable to annotate it with `@retry(times=0)`.
351
+ This decorator is useful if this step may hang indefinitely.
330
352
 
331
- This can be used in conjunction with the `@catch` decorator. The `@catch`
332
- decorator will execute a no-op task after all retries have been exhausted,
333
- ensuring that the flow execution can continue.
353
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
354
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
355
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
356
+
357
+ Note that all the values specified in parameters are added together so if you specify
358
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
334
359
  """
335
360
  ...
336
361
 
337
362
  @typing.overload
338
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
363
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
339
364
  ...
340
365
 
341
366
  @typing.overload
342
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
367
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
343
368
  ...
344
369
 
345
- 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):
370
+ 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):
346
371
  """
347
- Specifies the number of times the task corresponding
348
- to a step needs to be retried.
372
+ Specifies a timeout for your step.
349
373
 
350
- This decorator is useful for handling transient errors, such as networking issues.
351
- If your task contains operations that can't be retried safely, e.g. database updates,
352
- it is advisable to annotate it with `@retry(times=0)`.
374
+ This decorator is useful if this step may hang indefinitely.
353
375
 
354
- This can be used in conjunction with the `@catch` decorator. The `@catch`
355
- decorator will execute a no-op task after all retries have been exhausted,
356
- ensuring that the flow execution can continue.
376
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
377
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
378
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
379
+
380
+ Note that all the values specified in parameters are added together so if you specify
381
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
357
382
  """
358
383
  ...
359
384
 
@@ -364,92 +389,116 @@ def huggingface_hub(*, temp_dir_root: typing.Optional[str] = None, load: typing.
364
389
  ...
365
390
 
366
391
  @typing.overload
367
- 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]]]:
392
+ 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]]]:
368
393
  """
369
- Specifies the Conda environment for the step.
394
+ Specifies the PyPI packages for the step.
370
395
 
371
396
  Information in this decorator will augment any
372
- attributes set in the `@conda_base` flow-level decorator. Hence,
373
- you can use `@conda_base` to set packages required by all
374
- steps and use `@conda` to specify step-specific overrides.
397
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
398
+ you can use `@pypi_base` to set packages required by all
399
+ steps and use `@pypi` to specify step-specific overrides.
375
400
  """
376
401
  ...
377
402
 
378
403
  @typing.overload
379
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
404
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
380
405
  ...
381
406
 
382
407
  @typing.overload
383
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
408
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
384
409
  ...
385
410
 
386
- 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):
411
+ 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):
387
412
  """
388
- Specifies the Conda environment for the step.
413
+ Specifies the PyPI packages for the step.
389
414
 
390
415
  Information in this decorator will augment any
391
- attributes set in the `@conda_base` flow-level decorator. Hence,
392
- you can use `@conda_base` to set packages required by all
393
- steps and use `@conda` to specify step-specific overrides.
416
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
417
+ you can use `@pypi_base` to set packages required by all
418
+ steps and use `@pypi` to specify step-specific overrides.
394
419
  """
395
420
  ...
396
421
 
397
422
  @typing.overload
398
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
423
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
399
424
  """
400
- Internal decorator to support Fast bakery
425
+ Decorator prototype for all step decorators. This function gets specialized
426
+ and imported for all decorators types by _import_plugin_decorators().
401
427
  """
402
428
  ...
403
429
 
404
430
  @typing.overload
405
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
431
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
406
432
  ...
407
433
 
408
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
434
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
409
435
  """
410
- Internal decorator to support Fast bakery
436
+ Decorator prototype for all step decorators. This function gets specialized
437
+ and imported for all decorators types by _import_plugin_decorators().
411
438
  """
412
439
  ...
413
440
 
414
441
  @typing.overload
415
- 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]]]:
442
+ 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]]]:
416
443
  """
417
- Enables loading / saving of models within a step.
444
+ Specifies the number of times the task corresponding
445
+ to a step needs to be retried.
446
+
447
+ This decorator is useful for handling transient errors, such as networking issues.
448
+ If your task contains operations that can't be retried safely, e.g. database updates,
449
+ it is advisable to annotate it with `@retry(times=0)`.
450
+
451
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
452
+ decorator will execute a no-op task after all retries have been exhausted,
453
+ ensuring that the flow execution can continue.
418
454
  """
419
455
  ...
420
456
 
421
457
  @typing.overload
422
- def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
458
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
423
459
  ...
424
460
 
425
461
  @typing.overload
426
- def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
462
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
427
463
  ...
428
464
 
429
- 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):
465
+ def retry(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, times: int = 3, minutes_between_retries: int = 2):
430
466
  """
431
- Enables loading / saving of models within a step.
467
+ Specifies the number of times the task corresponding
468
+ to a step needs to be retried.
469
+
470
+ This decorator is useful for handling transient errors, such as networking issues.
471
+ If your task contains operations that can't be retried safely, e.g. database updates,
472
+ it is advisable to annotate it with `@retry(times=0)`.
473
+
474
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
475
+ decorator will execute a no-op task after all retries have been exhausted,
476
+ ensuring that the flow execution can continue.
432
477
  """
433
478
  ...
434
479
 
435
480
  @typing.overload
436
- 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]]]:
481
+ 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]]]:
437
482
  """
438
- Specifies environment variables to be set prior to the execution of a step.
483
+ Creates a human-readable report, a Metaflow Card, after this step completes.
484
+
485
+ Note that you may add multiple `@card` decorators in a step with different parameters.
439
486
  """
440
487
  ...
441
488
 
442
489
  @typing.overload
443
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
490
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
444
491
  ...
445
492
 
446
493
  @typing.overload
447
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
494
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
448
495
  ...
449
496
 
450
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
497
+ 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):
451
498
  """
452
- Specifies environment variables to be set prior to the execution of a step.
499
+ Creates a human-readable report, a Metaflow Card, after this step completes.
500
+
501
+ Note that you may add multiple `@card` decorators in a step with different parameters.
453
502
  """
454
503
  ...
455
504
 
@@ -476,67 +525,58 @@ def secrets(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None],
476
525
  """
477
526
  ...
478
527
 
479
- @typing.overload
480
- def resources(*, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
528
+ 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]]]:
481
529
  """
482
- Specifies the resources needed when executing this step.
483
-
484
- Use `@resources` to specify the resource requirements
485
- independently of the specific compute layer (`@batch`, `@kubernetes`).
486
-
487
- You can choose the compute layer on the command line by executing e.g.
488
- ```
489
- python myflow.py run --with batch
490
- ```
491
- or
492
- ```
493
- python myflow.py run --with kubernetes
494
- ```
495
- which executes the flow on the desired system using the
496
- requirements specified in `@resources`.
530
+ Specifies that this step should execute on Kubernetes.
531
+ """
532
+ ...
533
+
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]]:
535
+ """
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.
497
538
  """
498
539
  ...
499
540
 
500
541
  @typing.overload
501
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
542
+ 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]]:
543
+ """
544
+ Specifies the times when the flow should be run when running on a
545
+ production scheduler.
546
+ """
502
547
  ...
503
548
 
504
549
  @typing.overload
505
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
550
+ def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
506
551
  ...
507
552
 
508
- def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None):
553
+ 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):
509
554
  """
510
- Specifies the resources needed when executing this step.
511
-
512
- Use `@resources` to specify the resource requirements
513
- independently of the specific compute layer (`@batch`, `@kubernetes`).
514
-
515
- You can choose the compute layer on the command line by executing e.g.
516
- ```
517
- python myflow.py run --with batch
518
- ```
519
- or
520
- ```
521
- python myflow.py run --with kubernetes
522
- ```
523
- which executes the flow on the desired system using the
524
- requirements specified in `@resources`.
555
+ Specifies the times when the flow should be run when running on a
556
+ production scheduler.
525
557
  """
526
558
  ...
527
559
 
528
- 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]]]:
560
+ @typing.overload
561
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
529
562
  """
530
- Specifies that this step should execute on Kubernetes.
563
+ Specifies the PyPI packages for all steps of the flow.
564
+
565
+ Use `@pypi_base` to set common packages required by all
566
+ steps and use `@pypi` to specify step-specific overrides.
531
567
  """
532
568
  ...
533
569
 
534
- def project(*, name: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
570
+ @typing.overload
571
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
572
+ ...
573
+
574
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
535
575
  """
536
- Specifies what flows belong to the same project.
576
+ Specifies the PyPI packages for all steps of the flow.
537
577
 
538
- A project-specific namespace is created for all flows that
539
- use the same `@project(name)`.
578
+ Use `@pypi_base` to set common packages required by all
579
+ steps and use `@pypi` to specify step-specific overrides.
540
580
  """
541
581
  ...
542
582
 
@@ -550,6 +590,30 @@ def airflow_s3_key_sensor(*, timeout: int, poke_interval: int, mode: str, expone
550
590
  """
551
591
  ...
552
592
 
593
+ def nim(*, models: "list[NIM]", backend: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
594
+ """
595
+ This decorator is used to run NIM containers in Metaflow tasks as sidecars.
596
+
597
+ User code call
598
+ -----------
599
+ @nim(
600
+ models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
601
+ backend='managed'
602
+ )
603
+
604
+ Valid backend options
605
+ ---------------------
606
+ - 'managed': Outerbounds selects a compute provider based on the model.
607
+ - 🚧 'dataplane': Run in your account.
608
+
609
+ Valid model options
610
+ ----------------
611
+ - 'meta/llama3-8b-instruct': 8B parameter model
612
+ - 'meta/llama3-70b-instruct': 70B parameter model
613
+ - Upon request, any model here: https://nvcf.ngc.nvidia.com/functions?filter=nvidia-functions
614
+ """
615
+ ...
616
+
553
617
  @typing.overload
554
618
  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]]:
555
619
  """
@@ -631,55 +695,6 @@ def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *
631
695
  """
632
696
  ...
633
697
 
634
- 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]]:
635
- """
636
- 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.
637
- 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.
638
- """
639
- ...
640
-
641
- @typing.overload
642
- 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]]:
643
- """
644
- Specifies the times when the flow should be run when running on a
645
- production scheduler.
646
- """
647
- ...
648
-
649
- @typing.overload
650
- def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
651
- ...
652
-
653
- 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):
654
- """
655
- Specifies the times when the flow should be run when running on a
656
- production scheduler.
657
- """
658
- ...
659
-
660
- @typing.overload
661
- def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
662
- """
663
- Specifies the PyPI packages for all steps of the flow.
664
-
665
- Use `@pypi_base` to set common packages required by all
666
- steps and use `@pypi` to specify step-specific overrides.
667
- """
668
- ...
669
-
670
- @typing.overload
671
- def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
672
- ...
673
-
674
- def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
675
- """
676
- Specifies the PyPI packages for all steps of the flow.
677
-
678
- Use `@pypi_base` to set common packages required by all
679
- steps and use `@pypi` to specify step-specific overrides.
680
- """
681
- ...
682
-
683
698
  @typing.overload
684
699
  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]]:
685
700
  """
@@ -753,27 +768,12 @@ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: t
753
768
  """
754
769
  ...
755
770
 
756
- def nim(*, models: "list[NIM]", backend: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
771
+ def project(*, name: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
757
772
  """
758
- This decorator is used to run NIM containers in Metaflow tasks as sidecars.
759
-
760
- User code call
761
- -----------
762
- @nim(
763
- models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
764
- backend='managed'
765
- )
766
-
767
- Valid backend options
768
- ---------------------
769
- - 'managed': Outerbounds selects a compute provider based on the model.
770
- - 🚧 'dataplane': Run in your account.
773
+ Specifies what flows belong to the same project.
771
774
 
772
- Valid model options
773
- ----------------
774
- - 'meta/llama3-8b-instruct': 8B parameter model
775
- - 'meta/llama3-70b-instruct': 70B parameter model
776
- - Upon request, any model here: https://nvcf.ngc.nvidia.com/functions?filter=nvidia-functions
775
+ A project-specific namespace is created for all flows that
776
+ use the same `@project(name)`.
777
777
  """
778
778
  ...
779
779