ob-metaflow-stubs 6.0.3.150__py2.py3-none-any.whl → 6.0.3.152__py2.py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (198) hide show
  1. metaflow-stubs/__init__.pyi +599 -584
  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 +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 +2 -2
  20. metaflow-stubs/metaflow_current.pyi +140 -140
  21. metaflow-stubs/mf_extensions/__init__.pyi +2 -2
  22. metaflow-stubs/mf_extensions/obcheckpoint/__init__.pyi +2 -2
  23. metaflow-stubs/mf_extensions/obcheckpoint/plugins/__init__.pyi +2 -2
  24. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/__init__.pyi +2 -2
  25. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/__init__.pyi +2 -2
  26. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/async_cards.pyi +3 -3
  27. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/deco_injection_mixin.pyi +2 -2
  28. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/extra_components.pyi +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 +3 -3
  32. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/lineage_card.pyi +2 -2
  33. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/checkpoint_storage.pyi +4 -4
  34. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/constructors.pyi +2 -2
  35. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/core.pyi +3 -3
  36. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/decorator.pyi +3 -3
  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 +2 -2
  39. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/lineage.pyi +2 -2
  40. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/__init__.pyi +2 -2
  41. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/core.pyi +3 -3
  42. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/exceptions.pyi +2 -2
  43. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/task_utils.pyi +3 -3
  44. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/utils.pyi +2 -2
  45. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastructures.pyi +2 -2
  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 +2 -2
  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 +3 -3
  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 +2 -2
  60. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/tar_utils.pyi +2 -2
  61. metaflow-stubs/mf_extensions/outerbounds/__init__.pyi +2 -2
  62. metaflow-stubs/mf_extensions/outerbounds/plugins/__init__.pyi +2 -2
  63. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/__init__.pyi +2 -2
  64. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/injector.pyi +2 -2
  65. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/__init__.pyi +2 -2
  66. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/ollama.pyi +2 -2
  67. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/__init__.pyi +2 -2
  68. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/snowflake.pyi +2 -2
  69. metaflow-stubs/mf_extensions/outerbounds/profilers/__init__.pyi +2 -2
  70. metaflow-stubs/mf_extensions/outerbounds/profilers/gpu.pyi +2 -2
  71. metaflow-stubs/mf_extensions/outerbounds/remote_config.pyi +2 -2
  72. metaflow-stubs/mf_extensions/outerbounds/toplevel/__init__.pyi +2 -2
  73. metaflow-stubs/mf_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.pyi +2 -2
  74. metaflow-stubs/multicore_utils.pyi +2 -2
  75. metaflow-stubs/parameters.pyi +4 -4
  76. metaflow-stubs/plugins/__init__.pyi +12 -12
  77. metaflow-stubs/plugins/airflow/__init__.pyi +2 -2
  78. metaflow-stubs/plugins/airflow/airflow_utils.pyi +2 -2
  79. metaflow-stubs/plugins/airflow/exception.pyi +2 -2
  80. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +2 -2
  81. metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi +2 -2
  82. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +2 -2
  83. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +2 -2
  84. metaflow-stubs/plugins/argo/__init__.pyi +2 -2
  85. metaflow-stubs/plugins/argo/argo_client.pyi +2 -2
  86. metaflow-stubs/plugins/argo/argo_events.pyi +2 -2
  87. metaflow-stubs/plugins/argo/argo_workflows.pyi +3 -3
  88. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +4 -4
  89. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +4 -4
  90. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +2 -2
  91. metaflow-stubs/plugins/aws/__init__.pyi +2 -2
  92. metaflow-stubs/plugins/aws/aws_client.pyi +2 -2
  93. metaflow-stubs/plugins/aws/aws_utils.pyi +2 -2
  94. metaflow-stubs/plugins/aws/batch/__init__.pyi +2 -2
  95. metaflow-stubs/plugins/aws/batch/batch.pyi +2 -2
  96. metaflow-stubs/plugins/aws/batch/batch_client.pyi +2 -2
  97. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +2 -2
  98. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +2 -2
  99. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +3 -3
  100. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +2 -2
  101. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +2 -2
  102. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +2 -2
  103. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +2 -2
  104. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +2 -2
  105. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer.pyi +3 -3
  106. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer_objects.pyi +2 -2
  107. metaflow-stubs/plugins/azure/__init__.pyi +2 -2
  108. metaflow-stubs/plugins/azure/azure_credential.pyi +2 -2
  109. metaflow-stubs/plugins/azure/azure_exceptions.pyi +2 -2
  110. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +3 -3
  111. metaflow-stubs/plugins/azure/azure_utils.pyi +2 -2
  112. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +2 -2
  113. metaflow-stubs/plugins/azure/includefile_support.pyi +2 -2
  114. metaflow-stubs/plugins/cards/__init__.pyi +6 -6
  115. metaflow-stubs/plugins/cards/card_client.pyi +2 -2
  116. metaflow-stubs/plugins/cards/card_creator.pyi +2 -2
  117. metaflow-stubs/plugins/cards/card_datastore.pyi +2 -2
  118. metaflow-stubs/plugins/cards/card_decorator.pyi +2 -2
  119. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +2 -2
  120. metaflow-stubs/plugins/cards/card_modules/basic.pyi +2 -2
  121. metaflow-stubs/plugins/cards/card_modules/card.pyi +2 -2
  122. metaflow-stubs/plugins/cards/card_modules/components.pyi +4 -4
  123. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +2 -2
  124. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +2 -2
  125. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +2 -2
  126. metaflow-stubs/plugins/cards/card_resolver.pyi +2 -2
  127. metaflow-stubs/plugins/cards/component_serializer.pyi +2 -2
  128. metaflow-stubs/plugins/cards/exception.pyi +2 -2
  129. metaflow-stubs/plugins/catch_decorator.pyi +3 -3
  130. metaflow-stubs/plugins/datatools/__init__.pyi +2 -2
  131. metaflow-stubs/plugins/datatools/local.pyi +2 -2
  132. metaflow-stubs/plugins/datatools/s3/__init__.pyi +2 -2
  133. metaflow-stubs/plugins/datatools/s3/s3.pyi +2 -2
  134. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +2 -2
  135. metaflow-stubs/plugins/datatools/s3/s3util.pyi +2 -2
  136. metaflow-stubs/plugins/debug_logger.pyi +2 -2
  137. metaflow-stubs/plugins/debug_monitor.pyi +2 -2
  138. metaflow-stubs/plugins/environment_decorator.pyi +2 -2
  139. metaflow-stubs/plugins/events_decorator.pyi +2 -2
  140. metaflow-stubs/plugins/frameworks/__init__.pyi +2 -2
  141. metaflow-stubs/plugins/frameworks/pytorch.pyi +2 -2
  142. metaflow-stubs/plugins/gcp/__init__.pyi +2 -2
  143. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +3 -3
  144. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +2 -2
  145. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +2 -2
  146. metaflow-stubs/plugins/gcp/gs_utils.pyi +2 -2
  147. metaflow-stubs/plugins/gcp/includefile_support.pyi +2 -2
  148. metaflow-stubs/plugins/kubernetes/__init__.pyi +2 -2
  149. metaflow-stubs/plugins/kubernetes/kube_utils.pyi +2 -2
  150. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +2 -2
  151. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +2 -2
  152. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +2 -2
  153. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +2 -2
  154. metaflow-stubs/plugins/kubernetes/spot_monitor_sidecar.pyi +2 -2
  155. metaflow-stubs/plugins/ollama/__init__.pyi +3 -3
  156. metaflow-stubs/plugins/parallel_decorator.pyi +2 -2
  157. metaflow-stubs/plugins/perimeters.pyi +2 -2
  158. metaflow-stubs/plugins/project_decorator.pyi +2 -2
  159. metaflow-stubs/plugins/pypi/__init__.pyi +2 -2
  160. metaflow-stubs/plugins/pypi/conda_decorator.pyi +2 -2
  161. metaflow-stubs/plugins/pypi/conda_environment.pyi +3 -3
  162. metaflow-stubs/plugins/pypi/parsers.pyi +2 -2
  163. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +2 -2
  164. metaflow-stubs/plugins/pypi/pypi_environment.pyi +2 -2
  165. metaflow-stubs/plugins/pypi/utils.pyi +2 -2
  166. metaflow-stubs/plugins/resources_decorator.pyi +2 -2
  167. metaflow-stubs/plugins/retry_decorator.pyi +2 -2
  168. metaflow-stubs/plugins/secrets/__init__.pyi +2 -2
  169. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +3 -3
  170. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +2 -2
  171. metaflow-stubs/plugins/snowflake/__init__.pyi +2 -2
  172. metaflow-stubs/plugins/storage_executor.pyi +2 -2
  173. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +2 -2
  174. metaflow-stubs/plugins/timeout_decorator.pyi +3 -3
  175. metaflow-stubs/profilers/__init__.pyi +2 -2
  176. metaflow-stubs/pylint_wrapper.pyi +2 -2
  177. metaflow-stubs/runner/__init__.pyi +2 -2
  178. metaflow-stubs/runner/deployer.pyi +4 -4
  179. metaflow-stubs/runner/deployer_impl.pyi +2 -2
  180. metaflow-stubs/runner/metaflow_runner.pyi +3 -3
  181. metaflow-stubs/runner/nbdeploy.pyi +2 -2
  182. metaflow-stubs/runner/nbrun.pyi +2 -2
  183. metaflow-stubs/runner/subprocess_manager.pyi +2 -2
  184. metaflow-stubs/runner/utils.pyi +4 -4
  185. metaflow-stubs/system/__init__.pyi +2 -2
  186. metaflow-stubs/system/system_logger.pyi +2 -2
  187. metaflow-stubs/system/system_monitor.pyi +2 -2
  188. metaflow-stubs/tagging_util.pyi +2 -2
  189. metaflow-stubs/tuple_util.pyi +2 -2
  190. metaflow-stubs/user_configs/__init__.pyi +2 -2
  191. metaflow-stubs/user_configs/config_decorators.pyi +6 -6
  192. metaflow-stubs/user_configs/config_options.pyi +3 -3
  193. metaflow-stubs/user_configs/config_parameters.pyi +6 -6
  194. {ob_metaflow_stubs-6.0.3.150.dist-info → ob_metaflow_stubs-6.0.3.152.dist-info}/METADATA +1 -1
  195. ob_metaflow_stubs-6.0.3.152.dist-info/RECORD +198 -0
  196. ob_metaflow_stubs-6.0.3.150.dist-info/RECORD +0 -198
  197. {ob_metaflow_stubs-6.0.3.150.dist-info → ob_metaflow_stubs-6.0.3.152.dist-info}/WHEEL +0 -0
  198. {ob_metaflow_stubs-6.0.3.150.dist-info → ob_metaflow_stubs-6.0.3.152.dist-info}/top_level.txt +0 -0
@@ -1,7 +1,7 @@
1
1
  ######################################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
- # MF version: 2.15.5.1+obcheckpoint(0.1.9);ob(v1) #
4
- # Generated on 2025-03-13T17:31:16.313988 #
3
+ # MF version: 2.15.7.1+obcheckpoint(0.1.9);ob(v1) #
4
+ # Generated on 2025-03-24T18:51:07.966816 #
5
5
  ######################################################################################################
6
6
 
7
7
  from __future__ import annotations
@@ -150,51 +150,76 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
150
150
  ...
151
151
 
152
152
  @typing.overload
153
- 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]]]:
153
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
154
154
  """
155
- Creates a human-readable report, a Metaflow Card, after this step completes.
155
+ Internal decorator to support Fast bakery
156
+ """
157
+ ...
158
+
159
+ @typing.overload
160
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
161
+ ...
162
+
163
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
164
+ """
165
+ Internal decorator to support Fast bakery
166
+ """
167
+ ...
168
+
169
+ @typing.overload
170
+ 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]]]:
171
+ """
172
+ Enables loading / saving of models within a step.
156
173
 
157
- Note that you may add multiple `@card` decorators in a step with different parameters.
158
174
 
159
175
 
160
176
  Parameters
161
177
  ----------
162
- type : str, default 'default'
163
- Card type.
164
- id : str, optional, default None
165
- If multiple cards are present, use this id to identify this card.
166
- options : Dict[str, Any], default {}
167
- Options passed to the card. The contents depend on the card type.
168
- timeout : int, default 45
169
- Interrupt reporting if it takes more than this many seconds.
178
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
179
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
180
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
181
+ - `current.checkpoint`
182
+ - `current.model`
183
+ - `current.huggingface_hub`
184
+
185
+ If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
186
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
187
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
188
+
189
+ temp_dir_root : str, default: None
190
+ The root directory under which `current.model.loaded` will store loaded models
170
191
  """
171
192
  ...
172
193
 
173
194
  @typing.overload
174
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
195
+ def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
175
196
  ...
176
197
 
177
198
  @typing.overload
178
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
199
+ def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
179
200
  ...
180
201
 
181
- 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):
202
+ 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):
182
203
  """
183
- Creates a human-readable report, a Metaflow Card, after this step completes.
204
+ Enables loading / saving of models within a step.
184
205
 
185
- Note that you may add multiple `@card` decorators in a step with different parameters.
186
206
 
187
207
 
188
208
  Parameters
189
209
  ----------
190
- type : str, default 'default'
191
- Card type.
192
- id : str, optional, default None
193
- If multiple cards are present, use this id to identify this card.
194
- options : Dict[str, Any], default {}
195
- Options passed to the card. The contents depend on the card type.
196
- timeout : int, default 45
197
- Interrupt reporting if it takes more than this many seconds.
210
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
211
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
212
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
213
+ - `current.checkpoint`
214
+ - `current.model`
215
+ - `current.huggingface_hub`
216
+
217
+ If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
218
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
219
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
220
+
221
+ temp_dir_root : str, default: None
222
+ The root directory under which `current.model.loaded` will store loaded models
198
223
  """
199
224
  ...
200
225
 
@@ -258,78 +283,145 @@ def conda(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
258
283
  ...
259
284
 
260
285
  @typing.overload
261
- 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]]]:
286
+ def catch(*, var: typing.Optional[str] = None, print_exception: bool = True) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
262
287
  """
263
- Enables checkpointing for a step.
288
+ Specifies that the step will success under all circumstances.
264
289
 
290
+ The decorator will create an optional artifact, specified by `var`, which
291
+ contains the exception raised. You can use it to detect the presence
292
+ of errors, indicating that all happy-path artifacts produced by the step
293
+ are missing.
265
294
 
266
295
 
267
296
  Parameters
268
297
  ----------
269
- load_policy : str, default: "fresh"
270
- The policy for loading the checkpoint. The following policies are supported:
271
- - "eager": Loads the the latest available checkpoint within the namespace.
272
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
273
- will be loaded at the start of the task.
274
- - "none": Do not load any checkpoint
275
- - "fresh": Loads the lastest checkpoint created within the running Task.
276
- This mode helps loading checkpoints across various retry attempts of the same task.
277
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
278
- created within the task will be loaded when the task is retries execution on failure.
279
-
280
- temp_dir_root : str, default: None
281
- The root directory under which `current.checkpoint.directory` will be created.
298
+ var : str, optional, default None
299
+ Name of the artifact in which to store the caught exception.
300
+ If not specified, the exception is not stored.
301
+ print_exception : bool, default True
302
+ Determines whether or not the exception is printed to
303
+ stdout when caught.
282
304
  """
283
305
  ...
284
306
 
285
307
  @typing.overload
286
- def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
308
+ def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
287
309
  ...
288
310
 
289
311
  @typing.overload
290
- def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
312
+ def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
291
313
  ...
292
314
 
293
- 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):
315
+ def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, var: typing.Optional[str] = None, print_exception: bool = True):
294
316
  """
295
- Enables checkpointing for a step.
317
+ Specifies that the step will success under all circumstances.
296
318
 
319
+ The decorator will create an optional artifact, specified by `var`, which
320
+ contains the exception raised. You can use it to detect the presence
321
+ of errors, indicating that all happy-path artifacts produced by the step
322
+ are missing.
297
323
 
298
324
 
299
325
  Parameters
300
326
  ----------
301
- load_policy : str, default: "fresh"
302
- The policy for loading the checkpoint. The following policies are supported:
303
- - "eager": Loads the the latest available checkpoint within the namespace.
304
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
305
- will be loaded at the start of the task.
306
- - "none": Do not load any checkpoint
307
- - "fresh": Loads the lastest checkpoint created within the running Task.
308
- This mode helps loading checkpoints across various retry attempts of the same task.
309
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
310
- created within the task will be loaded when the task is retries execution on failure.
327
+ var : str, optional, default None
328
+ Name of the artifact in which to store the caught exception.
329
+ If not specified, the exception is not stored.
330
+ print_exception : bool, default True
331
+ Determines whether or not the exception is printed to
332
+ stdout when caught.
333
+ """
334
+ ...
335
+
336
+ 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]]]:
337
+ """
338
+ Decorator that helps cache, version and store models/datasets from huggingface hub.
311
339
 
312
- temp_dir_root : str, default: None
313
- The root directory under which `current.checkpoint.directory` will be created.
340
+
341
+ Parameters
342
+ ----------
343
+ temp_dir_root : str, optional
344
+ The root directory that will hold the temporary directory where objects will be downloaded.
345
+
346
+ load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
347
+ The list of repos (models/datasets) to load.
348
+
349
+ Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
350
+
351
+ - If repo (model/dataset) is not found in the datastore:
352
+ - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
353
+ - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
354
+ - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
355
+
356
+ - If repo is found in the datastore:
357
+ - Loads it directly from datastore to local path (can be temporary directory or specified path)
314
358
  """
315
359
  ...
316
360
 
317
361
  @typing.overload
318
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
362
+ 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]]]:
319
363
  """
320
- Decorator prototype for all step decorators. This function gets specialized
321
- and imported for all decorators types by _import_plugin_decorators().
364
+ Specifies the PyPI packages for the step.
365
+
366
+ Information in this decorator will augment any
367
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
368
+ you can use `@pypi_base` to set packages required by all
369
+ steps and use `@pypi` to specify step-specific overrides.
370
+
371
+
372
+ Parameters
373
+ ----------
374
+ packages : Dict[str, str], default: {}
375
+ Packages to use for this step. The key is the name of the package
376
+ and the value is the version to use.
377
+ python : str, optional, default: None
378
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
379
+ that the version used will correspond to the version of the Python interpreter used to start the run.
322
380
  """
323
381
  ...
324
382
 
325
383
  @typing.overload
326
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
384
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
327
385
  ...
328
386
 
329
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
387
+ @typing.overload
388
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
389
+ ...
390
+
391
+ 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):
330
392
  """
331
- Decorator prototype for all step decorators. This function gets specialized
332
- and imported for all decorators types by _import_plugin_decorators().
393
+ Specifies the PyPI packages for the step.
394
+
395
+ Information in this decorator will augment any
396
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
397
+ you can use `@pypi_base` to set packages required by all
398
+ steps and use `@pypi` to specify step-specific overrides.
399
+
400
+
401
+ Parameters
402
+ ----------
403
+ packages : Dict[str, str], default: {}
404
+ Packages to use for this step. The key is the name of the package
405
+ and the value is the version to use.
406
+ python : str, optional, default: None
407
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
408
+ that the version used will correspond to the version of the Python interpreter used to start the run.
409
+ """
410
+ ...
411
+
412
+ 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]]]:
413
+ """
414
+ Specifies that this step should execute on DGX cloud.
415
+
416
+
417
+ Parameters
418
+ ----------
419
+ gpu : int
420
+ Number of GPUs to use.
421
+ gpu_type : str
422
+ Type of Nvidia GPU to use.
423
+ queue_timeout : int
424
+ Time to keep the job in NVCF's queue.
333
425
  """
334
426
  ...
335
427
 
@@ -422,311 +514,43 @@ def retry(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
422
514
  ...
423
515
 
424
516
  @typing.overload
425
- 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]]]:
517
+ 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]]]:
426
518
  """
427
- Specifies the PyPI packages for the step.
519
+ Specifies the resources needed when executing this step.
428
520
 
429
- Information in this decorator will augment any
430
- attributes set in the `@pyi_base` flow-level decorator. Hence,
431
- you can use `@pypi_base` to set packages required by all
432
- steps and use `@pypi` to specify step-specific overrides.
521
+ Use `@resources` to specify the resource requirements
522
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
523
+
524
+ You can choose the compute layer on the command line by executing e.g.
525
+ ```
526
+ python myflow.py run --with batch
527
+ ```
528
+ or
529
+ ```
530
+ python myflow.py run --with kubernetes
531
+ ```
532
+ which executes the flow on the desired system using the
533
+ requirements specified in `@resources`.
433
534
 
434
535
 
435
536
  Parameters
436
537
  ----------
437
- packages : Dict[str, str], default: {}
438
- Packages to use for this step. The key is the name of the package
439
- and the value is the version to use.
440
- python : str, optional, default: None
441
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
442
- that the version used will correspond to the version of the Python interpreter used to start the run.
538
+ cpu : int, default 1
539
+ Number of CPUs required for this step.
540
+ gpu : int, optional, default None
541
+ Number of GPUs required for this step.
542
+ disk : int, optional, default None
543
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
544
+ memory : int, default 4096
545
+ Memory size (in MB) required for this step.
546
+ shared_memory : int, optional, default None
547
+ The value for the size (in MiB) of the /dev/shm volume for this step.
548
+ This parameter maps to the `--shm-size` option in Docker.
443
549
  """
444
550
  ...
445
551
 
446
552
  @typing.overload
447
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
448
- ...
449
-
450
- @typing.overload
451
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
452
- ...
453
-
454
- 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):
455
- """
456
- Specifies the PyPI packages for the step.
457
-
458
- Information in this decorator will augment any
459
- attributes set in the `@pyi_base` flow-level decorator. Hence,
460
- you can use `@pypi_base` to set packages required by all
461
- steps and use `@pypi` to specify step-specific overrides.
462
-
463
-
464
- Parameters
465
- ----------
466
- packages : Dict[str, str], default: {}
467
- Packages to use for this step. The key is the name of the package
468
- and the value is the version to use.
469
- python : str, optional, default: None
470
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
471
- that the version used will correspond to the version of the Python interpreter used to start the run.
472
- """
473
- ...
474
-
475
- @typing.overload
476
- def catch(*, var: typing.Optional[str] = None, print_exception: bool = True) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
477
- """
478
- Specifies that the step will success under all circumstances.
479
-
480
- The decorator will create an optional artifact, specified by `var`, which
481
- contains the exception raised. You can use it to detect the presence
482
- of errors, indicating that all happy-path artifacts produced by the step
483
- are missing.
484
-
485
-
486
- Parameters
487
- ----------
488
- var : str, optional, default None
489
- Name of the artifact in which to store the caught exception.
490
- If not specified, the exception is not stored.
491
- print_exception : bool, default True
492
- Determines whether or not the exception is printed to
493
- stdout when caught.
494
- """
495
- ...
496
-
497
- @typing.overload
498
- def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
499
- ...
500
-
501
- @typing.overload
502
- def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
503
- ...
504
-
505
- def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, var: typing.Optional[str] = None, print_exception: bool = True):
506
- """
507
- Specifies that the step will success under all circumstances.
508
-
509
- The decorator will create an optional artifact, specified by `var`, which
510
- contains the exception raised. You can use it to detect the presence
511
- of errors, indicating that all happy-path artifacts produced by the step
512
- are missing.
513
-
514
-
515
- Parameters
516
- ----------
517
- var : str, optional, default None
518
- Name of the artifact in which to store the caught exception.
519
- If not specified, the exception is not stored.
520
- print_exception : bool, default True
521
- Determines whether or not the exception is printed to
522
- stdout when caught.
523
- """
524
- ...
525
-
526
- 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]]]:
527
- """
528
- Decorator that helps cache, version and store models/datasets from huggingface hub.
529
-
530
-
531
- Parameters
532
- ----------
533
- temp_dir_root : str, optional
534
- The root directory that will hold the temporary directory where objects will be downloaded.
535
-
536
- load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
537
- The list of repos (models/datasets) to load.
538
-
539
- Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
540
-
541
- - If repo (model/dataset) is not found in the datastore:
542
- - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
543
- - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
544
- - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
545
-
546
- - If repo is found in the datastore:
547
- - Loads it directly from datastore to local path (can be temporary directory or specified path)
548
- """
549
- ...
550
-
551
- @typing.overload
552
- 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]]]:
553
- """
554
- Enables loading / saving of models within a step.
555
-
556
-
557
-
558
- Parameters
559
- ----------
560
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
561
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
562
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
563
- - `current.checkpoint`
564
- - `current.model`
565
- - `current.huggingface_hub`
566
-
567
- If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
568
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
569
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
570
-
571
- temp_dir_root : str, default: None
572
- The root directory under which `current.model.loaded` will store loaded models
573
- """
574
- ...
575
-
576
- @typing.overload
577
- def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
578
- ...
579
-
580
- @typing.overload
581
- def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
582
- ...
583
-
584
- 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):
585
- """
586
- Enables loading / saving of models within a step.
587
-
588
-
589
-
590
- Parameters
591
- ----------
592
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
593
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
594
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
595
- - `current.checkpoint`
596
- - `current.model`
597
- - `current.huggingface_hub`
598
-
599
- If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
600
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
601
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
602
-
603
- temp_dir_root : str, default: None
604
- The root directory under which `current.model.loaded` will store loaded models
605
- """
606
- ...
607
-
608
- @typing.overload
609
- 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]]]:
610
- """
611
- Specifies secrets to be retrieved and injected as environment variables prior to
612
- the execution of a step.
613
-
614
-
615
- Parameters
616
- ----------
617
- sources : List[Union[str, Dict[str, Any]]], default: []
618
- List of secret specs, defining how the secrets are to be retrieved
619
- """
620
- ...
621
-
622
- @typing.overload
623
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
624
- ...
625
-
626
- @typing.overload
627
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
628
- ...
629
-
630
- 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]]] = []):
631
- """
632
- Specifies secrets to be retrieved and injected as environment variables prior to
633
- the execution of a step.
634
-
635
-
636
- Parameters
637
- ----------
638
- sources : List[Union[str, Dict[str, Any]]], default: []
639
- List of secret specs, defining how the secrets are to be retrieved
640
- """
641
- ...
642
-
643
- 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]]]:
644
- """
645
- Specifies that this step should execute on DGX cloud.
646
-
647
-
648
- Parameters
649
- ----------
650
- gpu : int
651
- Number of GPUs to use.
652
- gpu_type : str
653
- Type of Nvidia GPU to use.
654
- queue_timeout : int
655
- Time to keep the job in NVCF's queue.
656
- """
657
- ...
658
-
659
- def ollama(*, models: "list[Ollama]", backend: str) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
660
- """
661
- This decorator is used to run Ollama APIs as Metaflow task sidecars.
662
-
663
- User code call
664
- -----------
665
- @ollama(
666
- models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
667
- backend='local'
668
- )
669
-
670
- Valid backend options
671
- ---------------------
672
- - 'local': Run as a separate process on the local task machine.
673
- - (TODO) 'managed': Outerbounds hosts and selects compute provider.
674
- - (TODO) 'remote': Spin up separate instance to serve Ollama models.
675
-
676
- Valid model options
677
- ----------------
678
- - 'llama3.2'
679
- - 'llama3.3'
680
- - any model here https://ollama.com/search
681
-
682
-
683
- Parameters
684
- ----------
685
- models: list[Ollama]
686
- List of Ollama containers running models in sidecars.
687
- backend: str
688
- Determines where and how to run the Ollama process.
689
- """
690
- ...
691
-
692
- @typing.overload
693
- 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]]]:
694
- """
695
- Specifies the resources needed when executing this step.
696
-
697
- Use `@resources` to specify the resource requirements
698
- independently of the specific compute layer (`@batch`, `@kubernetes`).
699
-
700
- You can choose the compute layer on the command line by executing e.g.
701
- ```
702
- python myflow.py run --with batch
703
- ```
704
- or
705
- ```
706
- python myflow.py run --with kubernetes
707
- ```
708
- which executes the flow on the desired system using the
709
- requirements specified in `@resources`.
710
-
711
-
712
- Parameters
713
- ----------
714
- cpu : int, default 1
715
- Number of CPUs required for this step.
716
- gpu : int, optional, default None
717
- Number of GPUs required for this step.
718
- disk : int, optional, default None
719
- Disk size (in MB) required for this step. Only applies on Kubernetes.
720
- memory : int, default 4096
721
- Memory size (in MB) required for this step.
722
- shared_memory : int, optional, default None
723
- The value for the size (in MiB) of the /dev/shm volume for this step.
724
- This parameter maps to the `--shm-size` option in Docker.
725
- """
726
- ...
727
-
728
- @typing.overload
729
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
553
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
730
554
  ...
731
555
 
732
556
  @typing.overload
@@ -734,146 +558,37 @@ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None])
734
558
  ...
735
559
 
736
560
  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):
737
- """
738
- Specifies the resources needed when executing this step.
739
-
740
- Use `@resources` to specify the resource requirements
741
- independently of the specific compute layer (`@batch`, `@kubernetes`).
742
-
743
- You can choose the compute layer on the command line by executing e.g.
744
- ```
745
- python myflow.py run --with batch
746
- ```
747
- or
748
- ```
749
- python myflow.py run --with kubernetes
750
- ```
751
- which executes the flow on the desired system using the
752
- requirements specified in `@resources`.
753
-
754
-
755
- Parameters
756
- ----------
757
- cpu : int, default 1
758
- Number of CPUs required for this step.
759
- gpu : int, optional, default None
760
- Number of GPUs required for this step.
761
- disk : int, optional, default None
762
- Disk size (in MB) required for this step. Only applies on Kubernetes.
763
- memory : int, default 4096
764
- Memory size (in MB) required for this step.
765
- shared_memory : int, optional, default None
766
- The value for the size (in MiB) of the /dev/shm volume for this step.
767
- This parameter maps to the `--shm-size` option in Docker.
768
- """
769
- ...
770
-
771
- @typing.overload
772
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
773
- """
774
- Internal decorator to support Fast bakery
775
- """
776
- ...
777
-
778
- @typing.overload
779
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
780
- ...
781
-
782
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
783
- """
784
- Internal decorator to support Fast bakery
785
- """
786
- ...
787
-
788
- @typing.overload
789
- def environment(*, vars: typing.Dict[str, str] = {}) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
790
- """
791
- Specifies environment variables to be set prior to the execution of a step.
792
-
793
-
794
- Parameters
795
- ----------
796
- vars : Dict[str, str], default {}
797
- Dictionary of environment variables to set.
798
- """
799
- ...
800
-
801
- @typing.overload
802
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
803
- ...
804
-
805
- @typing.overload
806
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
807
- ...
808
-
809
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
810
- """
811
- Specifies environment variables to be set prior to the execution of a step.
812
-
813
-
814
- Parameters
815
- ----------
816
- vars : Dict[str, str], default {}
817
- Dictionary of environment variables to set.
818
- """
819
- ...
820
-
821
- @typing.overload
822
- 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]]]:
823
- """
824
- Specifies a timeout for your step.
825
-
826
- This decorator is useful if this step may hang indefinitely.
827
-
828
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
829
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
830
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
831
-
832
- Note that all the values specified in parameters are added together so if you specify
833
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
834
-
835
-
836
- Parameters
837
- ----------
838
- seconds : int, default 0
839
- Number of seconds to wait prior to timing out.
840
- minutes : int, default 0
841
- Number of minutes to wait prior to timing out.
842
- hours : int, default 0
843
- Number of hours to wait prior to timing out.
844
- """
845
- ...
846
-
847
- @typing.overload
848
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
849
- ...
850
-
851
- @typing.overload
852
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
853
- ...
854
-
855
- 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):
856
- """
857
- Specifies a timeout for your step.
858
-
859
- This decorator is useful if this step may hang indefinitely.
561
+ """
562
+ Specifies the resources needed when executing this step.
860
563
 
861
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
862
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
863
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
564
+ Use `@resources` to specify the resource requirements
565
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
864
566
 
865
- Note that all the values specified in parameters are added together so if you specify
866
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
567
+ You can choose the compute layer on the command line by executing e.g.
568
+ ```
569
+ python myflow.py run --with batch
570
+ ```
571
+ or
572
+ ```
573
+ python myflow.py run --with kubernetes
574
+ ```
575
+ which executes the flow on the desired system using the
576
+ requirements specified in `@resources`.
867
577
 
868
578
 
869
579
  Parameters
870
580
  ----------
871
- seconds : int, default 0
872
- Number of seconds to wait prior to timing out.
873
- minutes : int, default 0
874
- Number of minutes to wait prior to timing out.
875
- hours : int, default 0
876
- Number of hours to wait prior to timing out.
581
+ cpu : int, default 1
582
+ Number of CPUs required for this step.
583
+ gpu : int, optional, default None
584
+ Number of GPUs required for this step.
585
+ disk : int, optional, default None
586
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
587
+ memory : int, default 4096
588
+ Memory size (in MB) required for this step.
589
+ shared_memory : int, optional, default None
590
+ The value for the size (in MiB) of the /dev/shm volume for this step.
591
+ This parameter maps to the `--shm-size` option in Docker.
877
592
  """
878
593
  ...
879
594
 
@@ -954,87 +669,338 @@ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: ty
954
669
  """
955
670
  ...
956
671
 
957
- 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]]:
672
+ @typing.overload
673
+ 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]]]:
674
+ """
675
+ Specifies a timeout for your step.
676
+
677
+ This decorator is useful if this step may hang indefinitely.
678
+
679
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
680
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
681
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
682
+
683
+ Note that all the values specified in parameters are added together so if you specify
684
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
685
+
686
+
687
+ Parameters
688
+ ----------
689
+ seconds : int, default 0
690
+ Number of seconds to wait prior to timing out.
691
+ minutes : int, default 0
692
+ Number of minutes to wait prior to timing out.
693
+ hours : int, default 0
694
+ Number of hours to wait prior to timing out.
695
+ """
696
+ ...
697
+
698
+ @typing.overload
699
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
700
+ ...
701
+
702
+ @typing.overload
703
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
704
+ ...
705
+
706
+ 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):
707
+ """
708
+ Specifies a timeout for your step.
709
+
710
+ This decorator is useful if this step may hang indefinitely.
711
+
712
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
713
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
714
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
715
+
716
+ Note that all the values specified in parameters are added together so if you specify
717
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
718
+
719
+
720
+ Parameters
721
+ ----------
722
+ seconds : int, default 0
723
+ Number of seconds to wait prior to timing out.
724
+ minutes : int, default 0
725
+ Number of minutes to wait prior to timing out.
726
+ hours : int, default 0
727
+ Number of hours to wait prior to timing out.
728
+ """
729
+ ...
730
+
731
+ @typing.overload
732
+ 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]]]:
733
+ """
734
+ Enables checkpointing for a step.
735
+
736
+
737
+
738
+ Parameters
739
+ ----------
740
+ load_policy : str, default: "fresh"
741
+ The policy for loading the checkpoint. The following policies are supported:
742
+ - "eager": Loads the the latest available checkpoint within the namespace.
743
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
744
+ will be loaded at the start of the task.
745
+ - "none": Do not load any checkpoint
746
+ - "fresh": Loads the lastest checkpoint created within the running Task.
747
+ This mode helps loading checkpoints across various retry attempts of the same task.
748
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
749
+ created within the task will be loaded when the task is retries execution on failure.
750
+
751
+ temp_dir_root : str, default: None
752
+ The root directory under which `current.checkpoint.directory` will be created.
753
+ """
754
+ ...
755
+
756
+ @typing.overload
757
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
758
+ ...
759
+
760
+ @typing.overload
761
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
762
+ ...
763
+
764
+ 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):
765
+ """
766
+ Enables checkpointing for a step.
767
+
768
+
769
+
770
+ Parameters
771
+ ----------
772
+ load_policy : str, default: "fresh"
773
+ The policy for loading the checkpoint. The following policies are supported:
774
+ - "eager": Loads the the latest available checkpoint within the namespace.
775
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
776
+ will be loaded at the start of the task.
777
+ - "none": Do not load any checkpoint
778
+ - "fresh": Loads the lastest checkpoint created within the running Task.
779
+ This mode helps loading checkpoints across various retry attempts of the same task.
780
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
781
+ created within the task will be loaded when the task is retries execution on failure.
782
+
783
+ temp_dir_root : str, default: None
784
+ The root directory under which `current.checkpoint.directory` will be created.
785
+ """
786
+ ...
787
+
788
+ @typing.overload
789
+ 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]]]:
790
+ """
791
+ Creates a human-readable report, a Metaflow Card, after this step completes.
792
+
793
+ Note that you may add multiple `@card` decorators in a step with different parameters.
794
+
795
+
796
+ Parameters
797
+ ----------
798
+ type : str, default 'default'
799
+ Card type.
800
+ id : str, optional, default None
801
+ If multiple cards are present, use this id to identify this card.
802
+ options : Dict[str, Any], default {}
803
+ Options passed to the card. The contents depend on the card type.
804
+ timeout : int, default 45
805
+ Interrupt reporting if it takes more than this many seconds.
806
+ """
807
+ ...
808
+
809
+ @typing.overload
810
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
811
+ ...
812
+
813
+ @typing.overload
814
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
815
+ ...
816
+
817
+ 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):
818
+ """
819
+ Creates a human-readable report, a Metaflow Card, after this step completes.
820
+
821
+ Note that you may add multiple `@card` decorators in a step with different parameters.
822
+
823
+
824
+ Parameters
825
+ ----------
826
+ type : str, default 'default'
827
+ Card type.
828
+ id : str, optional, default None
829
+ If multiple cards are present, use this id to identify this card.
830
+ options : Dict[str, Any], default {}
831
+ Options passed to the card. The contents depend on the card type.
832
+ timeout : int, default 45
833
+ Interrupt reporting if it takes more than this many seconds.
834
+ """
835
+ ...
836
+
837
+ def ollama(*, models: "list[Ollama]", backend: str) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
838
+ """
839
+ This decorator is used to run Ollama APIs as Metaflow task sidecars.
840
+
841
+ User code call
842
+ -----------
843
+ @ollama(
844
+ models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
845
+ backend='local'
846
+ )
847
+
848
+ Valid backend options
849
+ ---------------------
850
+ - 'local': Run as a separate process on the local task machine.
851
+ - (TODO) 'managed': Outerbounds hosts and selects compute provider.
852
+ - (TODO) 'remote': Spin up separate instance to serve Ollama models.
853
+
854
+ Valid model options
855
+ ----------------
856
+ - 'llama3.2'
857
+ - 'llama3.3'
858
+ - any model here https://ollama.com/search
859
+
860
+
861
+ Parameters
862
+ ----------
863
+ models: list[Ollama]
864
+ List of Ollama containers running models in sidecars.
865
+ backend: str
866
+ Determines where and how to run the Ollama process.
867
+ """
868
+ ...
869
+
870
+ @typing.overload
871
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
872
+ """
873
+ Decorator prototype for all step decorators. This function gets specialized
874
+ and imported for all decorators types by _import_plugin_decorators().
875
+ """
876
+ ...
877
+
878
+ @typing.overload
879
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
880
+ ...
881
+
882
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
883
+ """
884
+ Decorator prototype for all step decorators. This function gets specialized
885
+ and imported for all decorators types by _import_plugin_decorators().
886
+ """
887
+ ...
888
+
889
+ @typing.overload
890
+ 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]]]:
891
+ """
892
+ Specifies secrets to be retrieved and injected as environment variables prior to
893
+ the execution of a step.
894
+
895
+
896
+ Parameters
897
+ ----------
898
+ sources : List[Union[str, Dict[str, Any]]], default: []
899
+ List of secret specs, defining how the secrets are to be retrieved
900
+ """
901
+ ...
902
+
903
+ @typing.overload
904
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
905
+ ...
906
+
907
+ @typing.overload
908
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
909
+ ...
910
+
911
+ 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]]] = []):
958
912
  """
959
- 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)
960
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
961
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
962
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
963
- starts only after all sensors finish.
913
+ Specifies secrets to be retrieved and injected as environment variables prior to
914
+ the execution of a step.
964
915
 
965
916
 
966
917
  Parameters
967
918
  ----------
968
- timeout : int
969
- Time, in seconds before the task times out and fails. (Default: 3600)
970
- poke_interval : int
971
- Time in seconds that the job should wait in between each try. (Default: 60)
972
- mode : str
973
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
974
- exponential_backoff : bool
975
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
976
- pool : str
977
- the slot pool this task should run in,
978
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
979
- soft_fail : bool
980
- Set to true to mark the task as SKIPPED on failure. (Default: False)
981
- name : str
982
- Name of the sensor on Airflow
983
- description : str
984
- Description of sensor in the Airflow UI
985
- bucket_key : Union[str, List[str]]
986
- The key(s) being waited on. Supports full s3:// style url or relative path from root level.
987
- When it's specified as a full s3:// url, please leave `bucket_name` as None
988
- bucket_name : str
989
- Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
990
- When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
991
- wildcard_match : bool
992
- whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
993
- aws_conn_id : str
994
- a reference to the s3 connection on Airflow. (Default: None)
995
- verify : bool
996
- Whether or not to verify SSL certificates for S3 connection. (Default: None)
919
+ sources : List[Union[str, Dict[str, Any]]], default: []
920
+ List of secret specs, defining how the secrets are to be retrieved
997
921
  """
998
922
  ...
999
923
 
1000
924
  @typing.overload
1001
- def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
925
+ 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]]]:
1002
926
  """
1003
- Specifies the PyPI packages for all steps of the flow.
927
+ Specifies environment variables to be set prior to the execution of a step.
1004
928
 
1005
- Use `@pypi_base` to set common packages required by all
1006
- steps and use `@pypi` to specify step-specific overrides.
1007
929
 
1008
930
  Parameters
1009
931
  ----------
1010
- packages : Dict[str, str], default: {}
1011
- Packages to use for this flow. The key is the name of the package
1012
- and the value is the version to use.
1013
- python : str, optional, default: None
1014
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1015
- that the version used will correspond to the version of the Python interpreter used to start the run.
932
+ vars : Dict[str, str], default {}
933
+ Dictionary of environment variables to set.
1016
934
  """
1017
935
  ...
1018
936
 
1019
937
  @typing.overload
1020
- def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
938
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1021
939
  ...
1022
940
 
1023
- def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
941
+ @typing.overload
942
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
943
+ ...
944
+
945
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
1024
946
  """
1025
- Specifies the PyPI packages for all steps of the flow.
947
+ Specifies environment variables to be set prior to the execution of a step.
1026
948
 
1027
- Use `@pypi_base` to set common packages required by all
1028
- steps and use `@pypi` to specify step-specific overrides.
1029
949
 
1030
950
  Parameters
1031
951
  ----------
1032
- packages : Dict[str, str], default: {}
1033
- Packages to use for this flow. The key is the name of the package
1034
- and the value is the version to use.
1035
- python : str, optional, default: None
1036
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1037
- that the version used will correspond to the version of the Python interpreter used to start the run.
952
+ vars : Dict[str, str], default {}
953
+ Dictionary of environment variables to set.
954
+ """
955
+ ...
956
+
957
+ def app_deploy(*, app_port: int, app_name: str) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
958
+ """
959
+ Specifies that this step is used to deploy an instance of the app.
960
+ Requires that self.app_name, self.app_port, self.entrypoint and self.deployDir is set.
961
+
962
+
963
+ Parameters
964
+ ----------
965
+ app_port : int
966
+ Number of GPUs to use.
967
+ app_name : str
968
+ Name of the app to deploy.
969
+ """
970
+ ...
971
+
972
+ def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
973
+ """
974
+ Specifies what flows belong to the same project.
975
+
976
+ A project-specific namespace is created for all flows that
977
+ use the same `@project(name)`.
978
+
979
+
980
+ Parameters
981
+ ----------
982
+ name : str
983
+ Project name. Make sure that the name is unique amongst all
984
+ projects that use the same production scheduler. The name may
985
+ contain only lowercase alphanumeric characters and underscores.
986
+
987
+ branch : Optional[str], default None
988
+ The branch to use. If not specified, the branch is set to
989
+ `user.<username>` unless `production` is set to `True`. This can
990
+ also be set on the command line using `--branch` as a top-level option.
991
+ It is an error to specify `branch` in the decorator and on the command line.
992
+
993
+ production : bool, default False
994
+ Whether or not the branch is the production branch. This can also be set on the
995
+ command line using `--production` as a top-level option. It is an error to specify
996
+ `production` in the decorator and on the command line.
997
+ The project branch name will be:
998
+ - if `branch` is specified:
999
+ - if `production` is True: `prod.<branch>`
1000
+ - if `production` is False: `test.<branch>`
1001
+ - if `branch` is not specified:
1002
+ - if `production` is True: `prod`
1003
+ - if `production` is False: `user.<username>`
1038
1004
  """
1039
1005
  ...
1040
1006
 
@@ -1139,6 +1105,49 @@ def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *
1139
1105
  """
1140
1106
  ...
1141
1107
 
1108
+ 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]]:
1109
+ """
1110
+ 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)
1111
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1112
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1113
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
1114
+ starts only after all sensors finish.
1115
+
1116
+
1117
+ Parameters
1118
+ ----------
1119
+ timeout : int
1120
+ Time, in seconds before the task times out and fails. (Default: 3600)
1121
+ poke_interval : int
1122
+ Time in seconds that the job should wait in between each try. (Default: 60)
1123
+ mode : str
1124
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1125
+ exponential_backoff : bool
1126
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1127
+ pool : str
1128
+ the slot pool this task should run in,
1129
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1130
+ soft_fail : bool
1131
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1132
+ name : str
1133
+ Name of the sensor on Airflow
1134
+ description : str
1135
+ Description of sensor in the Airflow UI
1136
+ bucket_key : Union[str, List[str]]
1137
+ The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1138
+ When it's specified as a full s3:// url, please leave `bucket_name` as None
1139
+ bucket_name : str
1140
+ Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1141
+ When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1142
+ wildcard_match : bool
1143
+ whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1144
+ aws_conn_id : str
1145
+ a reference to the s3 connection on Airflow. (Default: None)
1146
+ verify : bool
1147
+ Whether or not to verify SSL certificates for S3 connection. (Default: None)
1148
+ """
1149
+ ...
1150
+
1142
1151
  @typing.overload
1143
1152
  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]]:
1144
1153
  """
@@ -1233,6 +1242,47 @@ def airflow_external_task_sensor(*, timeout: int, poke_interval: int, mode: str,
1233
1242
  """
1234
1243
  ...
1235
1244
 
1245
+ @typing.overload
1246
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1247
+ """
1248
+ Specifies the PyPI packages for all steps of the flow.
1249
+
1250
+ Use `@pypi_base` to set common packages required by all
1251
+ steps and use `@pypi` to specify step-specific overrides.
1252
+
1253
+ Parameters
1254
+ ----------
1255
+ packages : Dict[str, str], default: {}
1256
+ Packages to use for this flow. The key is the name of the package
1257
+ and the value is the version to use.
1258
+ python : str, optional, default: None
1259
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1260
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1261
+ """
1262
+ ...
1263
+
1264
+ @typing.overload
1265
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1266
+ ...
1267
+
1268
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1269
+ """
1270
+ Specifies the PyPI packages for all steps of the flow.
1271
+
1272
+ Use `@pypi_base` to set common packages required by all
1273
+ steps and use `@pypi` to specify step-specific overrides.
1274
+
1275
+ Parameters
1276
+ ----------
1277
+ packages : Dict[str, str], default: {}
1278
+ Packages to use for this flow. The key is the name of the package
1279
+ and the value is the version to use.
1280
+ python : str, optional, default: None
1281
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1282
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1283
+ """
1284
+ ...
1285
+
1236
1286
  @typing.overload
1237
1287
  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]]:
1238
1288
  """
@@ -1377,40 +1427,5 @@ def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly:
1377
1427
  """
1378
1428
  ...
1379
1429
 
1380
- def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1381
- """
1382
- Specifies what flows belong to the same project.
1383
-
1384
- A project-specific namespace is created for all flows that
1385
- use the same `@project(name)`.
1386
-
1387
-
1388
- Parameters
1389
- ----------
1390
- name : str
1391
- Project name. Make sure that the name is unique amongst all
1392
- projects that use the same production scheduler. The name may
1393
- contain only lowercase alphanumeric characters and underscores.
1394
-
1395
- branch : Optional[str], default None
1396
- The branch to use. If not specified, the branch is set to
1397
- `user.<username>` unless `production` is set to `True`. This can
1398
- also be set on the command line using `--branch` as a top-level option.
1399
- It is an error to specify `branch` in the decorator and on the command line.
1400
-
1401
- production : bool, default False
1402
- Whether or not the branch is the production branch. This can also be set on the
1403
- command line using `--production` as a top-level option. It is an error to specify
1404
- `production` in the decorator and on the command line.
1405
- The project branch name will be:
1406
- - if `branch` is specified:
1407
- - if `production` is True: `prod.<branch>`
1408
- - if `production` is False: `test.<branch>`
1409
- - if `branch` is not specified:
1410
- - if `production` is True: `prod`
1411
- - if `production` is False: `user.<username>`
1412
- """
1413
- ...
1414
-
1415
1430
  pkg_name: str
1416
1431