ob-metaflow-stubs 6.0.3.138__py2.py3-none-any.whl → 6.0.3.139__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 +282 -282
  2. metaflow-stubs/cards.pyi +2 -2
  3. metaflow-stubs/cli.pyi +2 -2
  4. metaflow-stubs/cli_components/__init__.pyi +2 -2
  5. metaflow-stubs/cli_components/utils.pyi +2 -2
  6. metaflow-stubs/client/__init__.pyi +2 -2
  7. metaflow-stubs/client/core.pyi +5 -5
  8. metaflow-stubs/client/filecache.pyi +2 -2
  9. metaflow-stubs/events.pyi +2 -2
  10. metaflow-stubs/exception.pyi +2 -2
  11. metaflow-stubs/flowspec.pyi +6 -6
  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 +2 -2
  20. metaflow-stubs/metaflow_current.pyi +55 -55
  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 +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 +6 -6
  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 +7 -4
  36. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/decorator.pyi +55 -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 +39 -14
  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 +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 +3 -3
  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 +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 +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 +4 -4
  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 +3 -3
  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 +3 -3
  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 +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 +6 -6
  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 +2 -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 +5 -5
  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 +5 -5
  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 +2 -2
  155. metaflow-stubs/plugins/pypi/conda_decorator.pyi +2 -2
  156. metaflow-stubs/plugins/pypi/conda_environment.pyi +6 -6
  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 +3 -3
  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 +31 -31
  173. metaflow-stubs/runner/deployer_impl.pyi +2 -2
  174. metaflow-stubs/runner/metaflow_runner.pyi +4 -4
  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 +3 -3
  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 +7 -7
  188. {ob_metaflow_stubs-6.0.3.138.dist-info → ob_metaflow_stubs-6.0.3.139.dist-info}/METADATA +1 -1
  189. ob_metaflow_stubs-6.0.3.139.dist-info/RECORD +192 -0
  190. ob_metaflow_stubs-6.0.3.138.dist-info/RECORD +0 -192
  191. {ob_metaflow_stubs-6.0.3.138.dist-info → ob_metaflow_stubs-6.0.3.139.dist-info}/WHEEL +0 -0
  192. {ob_metaflow_stubs-6.0.3.138.dist-info → ob_metaflow_stubs-6.0.3.139.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.9.2+obcheckpoint(0.1.7);ob(v1) #
4
- # Generated on 2025-02-03T21:17:39.682936 #
3
+ # MF version: 2.13.9.2+obcheckpoint(0.1.8);ob(v1) #
4
+ # Generated on 2025-02-07T01:44:39.667412 #
5
5
  ######################################################################################################
6
6
 
7
7
  from __future__ import annotations
8
8
 
9
9
  import typing
10
10
  if typing.TYPE_CHECKING:
11
- import datetime
12
11
  import typing
12
+ import datetime
13
13
  FlowSpecDerived = typing.TypeVar("FlowSpecDerived", bound="FlowSpec", contravariant=False, covariant=False)
14
14
  StepFlag = typing.NewType("StepFlag", bool)
15
15
 
@@ -35,8 +35,8 @@ from .user_configs.config_parameters import ConfigValue as ConfigValue
35
35
  from .user_configs.config_parameters import config_expr as config_expr
36
36
  from .user_configs.config_decorators import CustomFlowDecorator as CustomFlowDecorator
37
37
  from .user_configs.config_decorators import CustomStepDecorator as CustomStepDecorator
38
- from . import tuple_util as tuple_util
39
38
  from . import cards as cards
39
+ from . import tuple_util as tuple_util
40
40
  from . import events as events
41
41
  from . import runner as runner
42
42
  from . import plugins as plugins
@@ -147,72 +147,56 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
147
147
  ...
148
148
 
149
149
  @typing.overload
150
- 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]]]:
150
+ def secrets(*, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = []) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
151
151
  """
152
- Specifies the resources needed when executing this step.
153
-
154
- Use `@resources` to specify the resource requirements
155
- independently of the specific compute layer (`@batch`, `@kubernetes`).
156
-
157
- You can choose the compute layer on the command line by executing e.g.
158
- ```
159
- python myflow.py run --with batch
160
- ```
161
- or
162
- ```
163
- python myflow.py run --with kubernetes
164
- ```
165
- which executes the flow on the desired system using the
166
- requirements specified in `@resources`.
152
+ Specifies secrets to be retrieved and injected as environment variables prior to
153
+ the execution of a step.
167
154
  """
168
155
  ...
169
156
 
170
157
  @typing.overload
171
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
158
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
172
159
  ...
173
160
 
174
161
  @typing.overload
175
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
162
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
176
163
  ...
177
164
 
178
- 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):
165
+ def secrets(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = []):
179
166
  """
180
- Specifies the resources needed when executing this step.
181
-
182
- Use `@resources` to specify the resource requirements
183
- independently of the specific compute layer (`@batch`, `@kubernetes`).
184
-
185
- You can choose the compute layer on the command line by executing e.g.
186
- ```
187
- python myflow.py run --with batch
188
- ```
189
- or
190
- ```
191
- python myflow.py run --with kubernetes
192
- ```
193
- which executes the flow on the desired system using the
194
- requirements specified in `@resources`.
167
+ Specifies secrets to be retrieved and injected as environment variables prior to
168
+ the execution of a step.
195
169
  """
196
170
  ...
197
171
 
198
172
  @typing.overload
199
- 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]]]:
173
+ 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]]]:
200
174
  """
201
- Enables checkpointing for a step.
175
+ Specifies the PyPI packages for the step.
176
+
177
+ Information in this decorator will augment any
178
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
179
+ you can use `@pypi_base` to set packages required by all
180
+ steps and use `@pypi` to specify step-specific overrides.
202
181
  """
203
182
  ...
204
183
 
205
184
  @typing.overload
206
- def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
185
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
207
186
  ...
208
187
 
209
188
  @typing.overload
210
- def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
189
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
211
190
  ...
212
191
 
213
- 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):
192
+ 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):
214
193
  """
215
- Enables checkpointing for a step.
194
+ Specifies the PyPI packages for the step.
195
+
196
+ Information in this decorator will augment any
197
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
198
+ you can use `@pypi_base` to set packages required by all
199
+ steps and use `@pypi` to specify step-specific overrides.
216
200
  """
217
201
  ...
218
202
 
@@ -241,42 +225,48 @@ def card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typ
241
225
  """
242
226
  ...
243
227
 
228
+ def nvidia(*, gpu: int, gpu_type: str, queue_timeout: int) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
229
+ """
230
+ Specifies that this step should execute on DGX cloud.
231
+ """
232
+ ...
233
+
244
234
  @typing.overload
245
- 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]]]:
235
+ 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]]]:
246
236
  """
247
- Specifies the number of times the task corresponding
248
- to a step needs to be retried.
237
+ Specifies a timeout for your step.
249
238
 
250
- This decorator is useful for handling transient errors, such as networking issues.
251
- If your task contains operations that can't be retried safely, e.g. database updates,
252
- it is advisable to annotate it with `@retry(times=0)`.
239
+ This decorator is useful if this step may hang indefinitely.
253
240
 
254
- This can be used in conjunction with the `@catch` decorator. The `@catch`
255
- decorator will execute a no-op task after all retries have been exhausted,
256
- ensuring that the flow execution can continue.
241
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
242
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
243
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
244
+
245
+ Note that all the values specified in parameters are added together so if you specify
246
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
257
247
  """
258
248
  ...
259
249
 
260
250
  @typing.overload
261
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
251
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
262
252
  ...
263
253
 
264
254
  @typing.overload
265
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
255
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
266
256
  ...
267
257
 
268
- 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):
258
+ 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):
269
259
  """
270
- Specifies the number of times the task corresponding
271
- to a step needs to be retried.
260
+ Specifies a timeout for your step.
272
261
 
273
- This decorator is useful for handling transient errors, such as networking issues.
274
- If your task contains operations that can't be retried safely, e.g. database updates,
275
- it is advisable to annotate it with `@retry(times=0)`.
262
+ This decorator is useful if this step may hang indefinitely.
276
263
 
277
- This can be used in conjunction with the `@catch` decorator. The `@catch`
278
- decorator will execute a no-op task after all retries have been exhausted,
279
- ensuring that the flow execution can continue.
264
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
265
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
266
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
267
+
268
+ Note that all the values specified in parameters are added together so if you specify
269
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
280
270
  """
281
271
  ...
282
272
 
@@ -312,40 +302,41 @@ def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
312
302
  ...
313
303
 
314
304
  @typing.overload
315
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
316
- """
317
- Internal decorator to support Fast bakery
318
- """
319
- ...
320
-
321
- @typing.overload
322
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
323
- ...
324
-
325
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
326
- """
327
- Internal decorator to support Fast bakery
328
- """
329
- ...
330
-
331
- @typing.overload
332
- 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]]]:
305
+ 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]]]:
333
306
  """
334
- Enables loading / saving of models within a step.
307
+ Specifies the number of times the task corresponding
308
+ to a step needs to be retried.
309
+
310
+ This decorator is useful for handling transient errors, such as networking issues.
311
+ If your task contains operations that can't be retried safely, e.g. database updates,
312
+ it is advisable to annotate it with `@retry(times=0)`.
313
+
314
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
315
+ decorator will execute a no-op task after all retries have been exhausted,
316
+ ensuring that the flow execution can continue.
335
317
  """
336
318
  ...
337
319
 
338
320
  @typing.overload
339
- def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
321
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
340
322
  ...
341
323
 
342
324
  @typing.overload
343
- def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
325
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
344
326
  ...
345
327
 
346
- 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):
328
+ 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):
347
329
  """
348
- Enables loading / saving of models within a step.
330
+ Specifies the number of times the task corresponding
331
+ to a step needs to be retried.
332
+
333
+ This decorator is useful for handling transient errors, such as networking issues.
334
+ If your task contains operations that can't be retried safely, e.g. database updates,
335
+ it is advisable to annotate it with `@retry(times=0)`.
336
+
337
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
338
+ decorator will execute a no-op task after all retries have been exhausted,
339
+ ensuring that the flow execution can continue.
349
340
  """
350
341
  ...
351
342
 
@@ -374,129 +365,111 @@ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: ty
374
365
  """
375
366
  ...
376
367
 
377
- def nvidia(*, gpu: int, gpu_type: str, queue_timeout: int) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
378
- """
379
- Specifies that this step should execute on DGX cloud.
380
- """
381
- ...
382
-
383
368
  @typing.overload
384
- 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]]]:
369
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
385
370
  """
386
- Specifies the PyPI packages for the step.
387
-
388
- Information in this decorator will augment any
389
- attributes set in the `@pyi_base` flow-level decorator. Hence,
390
- you can use `@pypi_base` to set packages required by all
391
- steps and use `@pypi` to specify step-specific overrides.
371
+ Internal decorator to support Fast bakery
392
372
  """
393
373
  ...
394
374
 
395
375
  @typing.overload
396
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
397
- ...
398
-
399
- @typing.overload
400
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
401
- ...
402
-
403
- 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):
404
- """
405
- Specifies the PyPI packages for the step.
406
-
407
- Information in this decorator will augment any
408
- attributes set in the `@pyi_base` flow-level decorator. Hence,
409
- you can use `@pypi_base` to set packages required by all
410
- steps and use `@pypi` to specify step-specific overrides.
411
- """
376
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
412
377
  ...
413
378
 
414
- 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]]]:
379
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
415
380
  """
416
- Decorator that helps cache, version and store models/datasets from huggingface hub.
381
+ Internal decorator to support Fast bakery
417
382
  """
418
383
  ...
419
384
 
420
385
  @typing.overload
421
- 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]]]:
386
+ 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]]]:
422
387
  """
423
- Specifies a timeout for your step.
424
-
425
- This decorator is useful if this step may hang indefinitely.
388
+ Specifies the resources needed when executing this step.
426
389
 
427
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
428
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
429
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
390
+ Use `@resources` to specify the resource requirements
391
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
430
392
 
431
- Note that all the values specified in parameters are added together so if you specify
432
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
393
+ You can choose the compute layer on the command line by executing e.g.
394
+ ```
395
+ python myflow.py run --with batch
396
+ ```
397
+ or
398
+ ```
399
+ python myflow.py run --with kubernetes
400
+ ```
401
+ which executes the flow on the desired system using the
402
+ requirements specified in `@resources`.
433
403
  """
434
404
  ...
435
405
 
436
406
  @typing.overload
437
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
407
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
438
408
  ...
439
409
 
440
410
  @typing.overload
441
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
411
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
442
412
  ...
443
413
 
444
- 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):
414
+ 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):
445
415
  """
446
- Specifies a timeout for your step.
447
-
448
- This decorator is useful if this step may hang indefinitely.
416
+ Specifies the resources needed when executing this step.
449
417
 
450
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
451
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
452
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
418
+ Use `@resources` to specify the resource requirements
419
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
453
420
 
454
- Note that all the values specified in parameters are added together so if you specify
455
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
421
+ You can choose the compute layer on the command line by executing e.g.
422
+ ```
423
+ python myflow.py run --with batch
424
+ ```
425
+ or
426
+ ```
427
+ python myflow.py run --with kubernetes
428
+ ```
429
+ which executes the flow on the desired system using the
430
+ requirements specified in `@resources`.
456
431
  """
457
432
  ...
458
433
 
459
434
  @typing.overload
460
- def secrets(*, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = []) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
435
+ 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]]]:
461
436
  """
462
- Specifies secrets to be retrieved and injected as environment variables prior to
463
- the execution of a step.
437
+ Enables checkpointing for a step.
464
438
  """
465
439
  ...
466
440
 
467
441
  @typing.overload
468
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
442
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
469
443
  ...
470
444
 
471
445
  @typing.overload
472
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
446
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
473
447
  ...
474
448
 
475
- def secrets(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = []):
449
+ 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):
476
450
  """
477
- Specifies secrets to be retrieved and injected as environment variables prior to
478
- the execution of a step.
451
+ Enables checkpointing for a step.
479
452
  """
480
453
  ...
481
454
 
482
455
  @typing.overload
483
- 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]]]:
456
+ 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]]]:
484
457
  """
485
- Specifies environment variables to be set prior to the execution of a step.
458
+ Enables loading / saving of models within a step.
486
459
  """
487
460
  ...
488
461
 
489
462
  @typing.overload
490
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
463
+ def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
491
464
  ...
492
465
 
493
466
  @typing.overload
494
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
467
+ def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
495
468
  ...
496
469
 
497
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
470
+ 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):
498
471
  """
499
- Specifies environment variables to be set prior to the execution of a step.
472
+ Enables loading / saving of models within a step.
500
473
  """
501
474
  ...
502
475
 
@@ -531,13 +504,53 @@ def conda(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
531
504
  """
532
505
  ...
533
506
 
534
- 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]]:
507
+ @typing.overload
508
+ 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]]]:
535
509
  """
536
- 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)
537
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
538
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
539
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
540
- starts only after all sensors finish.
510
+ Specifies environment variables to be set prior to the execution of a step.
511
+ """
512
+ ...
513
+
514
+ @typing.overload
515
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
516
+ ...
517
+
518
+ @typing.overload
519
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
520
+ ...
521
+
522
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
523
+ """
524
+ Specifies environment variables to be set prior to the execution of a step.
525
+ """
526
+ ...
527
+
528
+ 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]]]:
529
+ """
530
+ Decorator that helps cache, version and store models/datasets from huggingface hub.
531
+ """
532
+ ...
533
+
534
+ @typing.overload
535
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
536
+ """
537
+ Specifies the PyPI packages for all steps of the flow.
538
+
539
+ Use `@pypi_base` to set common packages required by all
540
+ steps and use `@pypi` to specify step-specific overrides.
541
+ """
542
+ ...
543
+
544
+ @typing.overload
545
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
546
+ ...
547
+
548
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
549
+ """
550
+ Specifies the PyPI packages for all steps of the flow.
551
+
552
+ Use `@pypi_base` to set common packages required by all
553
+ steps and use `@pypi` to specify step-specific overrides.
541
554
  """
542
555
  ...
543
556
 
@@ -549,75 +562,83 @@ def airflow_external_task_sensor(*, timeout: int, poke_interval: int, mode: str,
549
562
  ...
550
563
 
551
564
  @typing.overload
552
- 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]]:
565
+ 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]]:
553
566
  """
554
- Specifies the event(s) that this flow depends on.
567
+ Specifies the flow(s) that this flow depends on.
555
568
 
556
569
  ```
557
- @trigger(event='foo')
570
+ @trigger_on_finish(flow='FooFlow')
558
571
  ```
559
572
  or
560
573
  ```
561
- @trigger(events=['foo', 'bar'])
574
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
562
575
  ```
576
+ This decorator respects the @project decorator and triggers the flow
577
+ when upstream runs within the same namespace complete successfully
563
578
 
564
- Additionally, you can specify the parameter mappings
565
- to map event payload to Metaflow parameters for the flow.
579
+ Additionally, you can specify project aware upstream flow dependencies
580
+ by specifying the fully qualified project_flow_name.
566
581
  ```
567
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
582
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
568
583
  ```
569
584
  or
570
585
  ```
571
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
572
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
586
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
573
587
  ```
574
588
 
575
- 'parameters' can also be a list of strings and tuples like so:
576
- ```
577
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
578
- ```
579
- This is equivalent to:
589
+ You can also specify just the project or project branch (other values will be
590
+ inferred from the current project or project branch):
580
591
  ```
581
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
592
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
582
593
  ```
594
+
595
+ Note that `branch` is typically one of:
596
+ - `prod`
597
+ - `user.bob`
598
+ - `test.my_experiment`
599
+ - `prod.staging`
583
600
  """
584
601
  ...
585
602
 
586
603
  @typing.overload
587
- def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
604
+ def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
588
605
  ...
589
606
 
590
- def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: typing.Union[str, typing.Dict[str, typing.Any], None] = None, events: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], options: typing.Dict[str, typing.Any] = {}):
607
+ 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] = {}):
591
608
  """
592
- Specifies the event(s) that this flow depends on.
609
+ Specifies the flow(s) that this flow depends on.
593
610
 
594
611
  ```
595
- @trigger(event='foo')
612
+ @trigger_on_finish(flow='FooFlow')
596
613
  ```
597
614
  or
598
615
  ```
599
- @trigger(events=['foo', 'bar'])
616
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
600
617
  ```
618
+ This decorator respects the @project decorator and triggers the flow
619
+ when upstream runs within the same namespace complete successfully
601
620
 
602
- Additionally, you can specify the parameter mappings
603
- to map event payload to Metaflow parameters for the flow.
621
+ Additionally, you can specify project aware upstream flow dependencies
622
+ by specifying the fully qualified project_flow_name.
604
623
  ```
605
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
624
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
606
625
  ```
607
626
  or
608
627
  ```
609
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
610
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
628
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
611
629
  ```
612
630
 
613
- 'parameters' can also be a list of strings and tuples like so:
614
- ```
615
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
616
- ```
617
- This is equivalent to:
631
+ You can also specify just the project or project branch (other values will be
632
+ inferred from the current project or project branch):
618
633
  ```
619
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
634
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
620
635
  ```
636
+
637
+ Note that `branch` is typically one of:
638
+ - `prod`
639
+ - `user.bob`
640
+ - `test.my_experiment`
641
+ - `prod.staging`
621
642
  """
622
643
  ...
623
644
 
@@ -640,6 +661,15 @@ def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly:
640
661
  """
641
662
  ...
642
663
 
664
+ def project(*, name: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
665
+ """
666
+ Specifies what flows belong to the same project.
667
+
668
+ A project-specific namespace is created for all flows that
669
+ use the same `@project(name)`.
670
+ """
671
+ ...
672
+
643
673
  def nim(*, models: "list[NIM]", backend: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
644
674
  """
645
675
  This decorator is used to run NIM containers in Metaflow tasks as sidecars.
@@ -664,139 +694,109 @@ def nim(*, models: "list[NIM]", backend: str) -> typing.Callable[[typing.Type[Fl
664
694
  """
665
695
  ...
666
696
 
667
- def project(*, name: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
668
- """
669
- Specifies what flows belong to the same project.
670
-
671
- A project-specific namespace is created for all flows that
672
- use the same `@project(name)`.
673
- """
674
- ...
675
-
676
- @typing.overload
677
- def conda_base(*, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
678
- """
679
- Specifies the Conda environment for all steps of the flow.
680
-
681
- Use `@conda_base` to set common libraries required by all
682
- steps and use `@conda` to specify step-specific additions.
683
- """
684
- ...
685
-
686
- @typing.overload
687
- def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
688
- ...
689
-
690
- def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
691
- """
692
- Specifies the Conda environment for all steps of the flow.
693
-
694
- Use `@conda_base` to set common libraries required by all
695
- steps and use `@conda` to specify step-specific additions.
696
- """
697
- ...
698
-
699
- @typing.overload
700
- def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
701
- """
702
- Specifies the PyPI packages for all steps of the flow.
703
-
704
- Use `@pypi_base` to set common packages required by all
705
- steps and use `@pypi` to specify step-specific overrides.
706
- """
707
- ...
708
-
709
697
  @typing.overload
710
- def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
711
- ...
712
-
713
- def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
714
- """
715
- Specifies the PyPI packages for all steps of the flow.
716
-
717
- Use `@pypi_base` to set common packages required by all
718
- steps and use `@pypi` to specify step-specific overrides.
719
- """
720
- ...
721
-
722
- @typing.overload
723
- 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]]:
698
+ 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]]:
724
699
  """
725
- Specifies the flow(s) that this flow depends on.
700
+ Specifies the event(s) that this flow depends on.
726
701
 
727
702
  ```
728
- @trigger_on_finish(flow='FooFlow')
703
+ @trigger(event='foo')
729
704
  ```
730
705
  or
731
706
  ```
732
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
707
+ @trigger(events=['foo', 'bar'])
733
708
  ```
734
- This decorator respects the @project decorator and triggers the flow
735
- when upstream runs within the same namespace complete successfully
736
709
 
737
- Additionally, you can specify project aware upstream flow dependencies
738
- by specifying the fully qualified project_flow_name.
710
+ Additionally, you can specify the parameter mappings
711
+ to map event payload to Metaflow parameters for the flow.
739
712
  ```
740
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
713
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
741
714
  ```
742
715
  or
743
716
  ```
744
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
717
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
718
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
745
719
  ```
746
720
 
747
- You can also specify just the project or project branch (other values will be
748
- inferred from the current project or project branch):
721
+ 'parameters' can also be a list of strings and tuples like so:
749
722
  ```
750
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
723
+ @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
724
+ ```
725
+ This is equivalent to:
726
+ ```
727
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
751
728
  ```
752
-
753
- Note that `branch` is typically one of:
754
- - `prod`
755
- - `user.bob`
756
- - `test.my_experiment`
757
- - `prod.staging`
758
729
  """
759
730
  ...
760
731
 
761
732
  @typing.overload
762
- def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
733
+ def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
763
734
  ...
764
735
 
765
- 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] = {}):
736
+ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: typing.Union[str, typing.Dict[str, typing.Any], None] = None, events: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], options: typing.Dict[str, typing.Any] = {}):
766
737
  """
767
- Specifies the flow(s) that this flow depends on.
738
+ Specifies the event(s) that this flow depends on.
768
739
 
769
740
  ```
770
- @trigger_on_finish(flow='FooFlow')
741
+ @trigger(event='foo')
771
742
  ```
772
743
  or
773
744
  ```
774
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
745
+ @trigger(events=['foo', 'bar'])
775
746
  ```
776
- This decorator respects the @project decorator and triggers the flow
777
- when upstream runs within the same namespace complete successfully
778
747
 
779
- Additionally, you can specify project aware upstream flow dependencies
780
- by specifying the fully qualified project_flow_name.
748
+ Additionally, you can specify the parameter mappings
749
+ to map event payload to Metaflow parameters for the flow.
781
750
  ```
782
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
751
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
783
752
  ```
784
753
  or
785
754
  ```
786
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
755
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
756
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
787
757
  ```
788
758
 
789
- You can also specify just the project or project branch (other values will be
790
- inferred from the current project or project branch):
759
+ 'parameters' can also be a list of strings and tuples like so:
791
760
  ```
792
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
761
+ @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
762
+ ```
763
+ This is equivalent to:
764
+ ```
765
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
793
766
  ```
767
+ """
768
+ ...
769
+
770
+ @typing.overload
771
+ def conda_base(*, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
772
+ """
773
+ Specifies the Conda environment for all steps of the flow.
794
774
 
795
- Note that `branch` is typically one of:
796
- - `prod`
797
- - `user.bob`
798
- - `test.my_experiment`
799
- - `prod.staging`
775
+ Use `@conda_base` to set common libraries required by all
776
+ steps and use `@conda` to specify step-specific additions.
777
+ """
778
+ ...
779
+
780
+ @typing.overload
781
+ def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
782
+ ...
783
+
784
+ def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
785
+ """
786
+ Specifies the Conda environment for all steps of the flow.
787
+
788
+ Use `@conda_base` to set common libraries required by all
789
+ steps and use `@conda` to specify step-specific additions.
790
+ """
791
+ ...
792
+
793
+ 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]]:
794
+ """
795
+ 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)
796
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
797
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
798
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
799
+ starts only after all sensors finish.
800
800
  """
801
801
  ...
802
802