ob-metaflow-stubs 6.0.3.118__py2.py3-none-any.whl → 6.0.3.120__py2.py3-none-any.whl

Sign up to get free protection for your applications and to get access to all the features.
Files changed (174) hide show
  1. metaflow-stubs/__init__.pyi +243 -243
  2. metaflow-stubs/cards.pyi +2 -2
  3. metaflow-stubs/cli.pyi +2 -2
  4. metaflow-stubs/client/__init__.pyi +2 -2
  5. metaflow-stubs/client/core.pyi +5 -5
  6. metaflow-stubs/client/filecache.pyi +3 -3
  7. metaflow-stubs/clone_util.pyi +2 -2
  8. metaflow-stubs/events.pyi +3 -3
  9. metaflow-stubs/exception.pyi +2 -2
  10. metaflow-stubs/flowspec.pyi +5 -5
  11. metaflow-stubs/generated_for.txt +1 -1
  12. metaflow-stubs/includefile.pyi +3 -3
  13. metaflow-stubs/info_file.pyi +2 -2
  14. metaflow-stubs/metadata_provider/__init__.pyi +2 -2
  15. metaflow-stubs/metadata_provider/heartbeat.pyi +2 -2
  16. metaflow-stubs/metadata_provider/metadata.pyi +2 -2
  17. metaflow-stubs/metadata_provider/util.pyi +2 -2
  18. metaflow-stubs/metaflow_config.pyi +2 -2
  19. metaflow-stubs/metaflow_current.pyi +91 -91
  20. metaflow-stubs/mf_extensions/__init__.pyi +2 -2
  21. metaflow-stubs/mf_extensions/obcheckpoint/__init__.pyi +2 -2
  22. metaflow-stubs/mf_extensions/obcheckpoint/plugins/__init__.pyi +2 -2
  23. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/__init__.pyi +2 -2
  24. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/__init__.pyi +2 -2
  25. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/checkpoint_storage.pyi +3 -3
  26. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/constructors.pyi +2 -2
  27. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/core.pyi +5 -5
  28. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/exceptions.pyi +2 -2
  29. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/final_api.pyi +3 -3
  30. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/__init__.pyi +2 -2
  31. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/core.pyi +2 -2
  32. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/exceptions.pyi +2 -2
  33. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/task_utils.pyi +2 -2
  34. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/utils.pyi +2 -2
  35. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastructures.pyi +2 -2
  36. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/exceptions.pyi +2 -2
  37. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/__init__.pyi +2 -2
  38. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/flowspec_utils.pyi +2 -2
  39. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/general.pyi +2 -2
  40. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/identity_utils.pyi +2 -2
  41. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/tar_utils.pyi +2 -2
  42. metaflow-stubs/mf_extensions/outerbounds/__init__.pyi +2 -2
  43. metaflow-stubs/mf_extensions/outerbounds/profilers/__init__.pyi +2 -2
  44. metaflow-stubs/mf_extensions/outerbounds/profilers/gpu.pyi +2 -2
  45. metaflow-stubs/mf_extensions/outerbounds/remote_config.pyi +2 -2
  46. metaflow-stubs/mf_extensions/outerbounds/toplevel/__init__.pyi +2 -2
  47. metaflow-stubs/mf_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.pyi +2 -2
  48. metaflow-stubs/mflog/__init__.pyi +2 -2
  49. metaflow-stubs/mflog/mflog.pyi +2 -2
  50. metaflow-stubs/multicore_utils.pyi +2 -2
  51. metaflow-stubs/parameters.pyi +3 -3
  52. metaflow-stubs/plugins/__init__.pyi +13 -13
  53. metaflow-stubs/plugins/airflow/__init__.pyi +2 -2
  54. metaflow-stubs/plugins/airflow/airflow_utils.pyi +2 -2
  55. metaflow-stubs/plugins/airflow/exception.pyi +2 -2
  56. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +2 -2
  57. metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi +2 -2
  58. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +2 -2
  59. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +2 -2
  60. metaflow-stubs/plugins/argo/__init__.pyi +2 -2
  61. metaflow-stubs/plugins/argo/argo_client.pyi +2 -2
  62. metaflow-stubs/plugins/argo/argo_events.pyi +2 -2
  63. metaflow-stubs/plugins/argo/argo_workflows.pyi +2 -2
  64. metaflow-stubs/plugins/argo/argo_workflows_cli.pyi +2 -2
  65. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +3 -3
  66. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +3 -3
  67. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +2 -2
  68. metaflow-stubs/plugins/aws/__init__.pyi +2 -2
  69. metaflow-stubs/plugins/aws/aws_client.pyi +2 -2
  70. metaflow-stubs/plugins/aws/aws_utils.pyi +2 -2
  71. metaflow-stubs/plugins/aws/batch/__init__.pyi +2 -2
  72. metaflow-stubs/plugins/aws/batch/batch.pyi +2 -2
  73. metaflow-stubs/plugins/aws/batch/batch_client.pyi +2 -2
  74. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +2 -2
  75. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +2 -2
  76. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +4 -4
  77. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +2 -2
  78. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +2 -2
  79. metaflow-stubs/plugins/aws/step_functions/production_token.pyi +2 -2
  80. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +2 -2
  81. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +2 -2
  82. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +2 -2
  83. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer.pyi +3 -3
  84. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer_objects.pyi +3 -3
  85. metaflow-stubs/plugins/azure/__init__.pyi +2 -2
  86. metaflow-stubs/plugins/azure/azure_credential.pyi +2 -2
  87. metaflow-stubs/plugins/azure/azure_exceptions.pyi +2 -2
  88. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +4 -4
  89. metaflow-stubs/plugins/azure/azure_utils.pyi +2 -2
  90. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +2 -2
  91. metaflow-stubs/plugins/azure/includefile_support.pyi +2 -2
  92. metaflow-stubs/plugins/cards/__init__.pyi +2 -2
  93. metaflow-stubs/plugins/cards/card_cli.pyi +2 -2
  94. metaflow-stubs/plugins/cards/card_client.pyi +2 -2
  95. metaflow-stubs/plugins/cards/card_creator.pyi +2 -2
  96. metaflow-stubs/plugins/cards/card_datastore.pyi +2 -2
  97. metaflow-stubs/plugins/cards/card_decorator.pyi +2 -2
  98. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +2 -2
  99. metaflow-stubs/plugins/cards/card_modules/basic.pyi +3 -3
  100. metaflow-stubs/plugins/cards/card_modules/card.pyi +2 -2
  101. metaflow-stubs/plugins/cards/card_modules/components.pyi +4 -4
  102. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +2 -2
  103. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +2 -2
  104. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +2 -2
  105. metaflow-stubs/plugins/cards/card_resolver.pyi +2 -2
  106. metaflow-stubs/plugins/cards/component_serializer.pyi +2 -2
  107. metaflow-stubs/plugins/cards/exception.pyi +2 -2
  108. metaflow-stubs/plugins/catch_decorator.pyi +3 -3
  109. metaflow-stubs/plugins/datatools/__init__.pyi +2 -2
  110. metaflow-stubs/plugins/datatools/local.pyi +2 -2
  111. metaflow-stubs/plugins/datatools/s3/__init__.pyi +2 -2
  112. metaflow-stubs/plugins/datatools/s3/s3.pyi +3 -3
  113. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +2 -2
  114. metaflow-stubs/plugins/datatools/s3/s3util.pyi +2 -2
  115. metaflow-stubs/plugins/debug_logger.pyi +2 -2
  116. metaflow-stubs/plugins/debug_monitor.pyi +2 -2
  117. metaflow-stubs/plugins/environment_decorator.pyi +2 -2
  118. metaflow-stubs/plugins/events_decorator.pyi +12 -2
  119. metaflow-stubs/plugins/frameworks/__init__.pyi +2 -2
  120. metaflow-stubs/plugins/frameworks/pytorch.pyi +2 -2
  121. metaflow-stubs/plugins/gcp/__init__.pyi +2 -2
  122. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +4 -4
  123. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +2 -2
  124. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +2 -2
  125. metaflow-stubs/plugins/gcp/gs_utils.pyi +2 -2
  126. metaflow-stubs/plugins/gcp/includefile_support.pyi +2 -2
  127. metaflow-stubs/plugins/kubernetes/__init__.pyi +3 -3
  128. metaflow-stubs/plugins/kubernetes/kube_utils.pyi +2 -2
  129. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +2 -2
  130. metaflow-stubs/plugins/kubernetes/kubernetes_cli.pyi +2 -2
  131. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +2 -2
  132. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +2 -2
  133. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +2 -2
  134. metaflow-stubs/plugins/logs_cli.pyi +2 -2
  135. metaflow-stubs/plugins/package_cli.pyi +2 -2
  136. metaflow-stubs/plugins/parallel_decorator.pyi +2 -2
  137. metaflow-stubs/plugins/perimeters.pyi +2 -2
  138. metaflow-stubs/plugins/project_decorator.pyi +2 -2
  139. metaflow-stubs/plugins/pypi/__init__.pyi +2 -2
  140. metaflow-stubs/plugins/pypi/conda_decorator.pyi +2 -2
  141. metaflow-stubs/plugins/pypi/conda_environment.pyi +4 -4
  142. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +2 -2
  143. metaflow-stubs/plugins/pypi/pypi_environment.pyi +2 -2
  144. metaflow-stubs/plugins/pypi/utils.pyi +2 -2
  145. metaflow-stubs/plugins/resources_decorator.pyi +2 -2
  146. metaflow-stubs/plugins/retry_decorator.pyi +2 -2
  147. metaflow-stubs/plugins/secrets/__init__.pyi +3 -3
  148. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +3 -3
  149. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +2 -2
  150. metaflow-stubs/plugins/storage_executor.pyi +2 -2
  151. metaflow-stubs/plugins/tag_cli.pyi +2 -2
  152. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +2 -2
  153. metaflow-stubs/plugins/timeout_decorator.pyi +3 -3
  154. metaflow-stubs/procpoll.pyi +2 -2
  155. metaflow-stubs/profilers/__init__.pyi +2 -2
  156. metaflow-stubs/pylint_wrapper.pyi +2 -2
  157. metaflow-stubs/runner/__init__.pyi +2 -2
  158. metaflow-stubs/runner/deployer.pyi +28 -28
  159. metaflow-stubs/runner/deployer_impl.pyi +2 -2
  160. metaflow-stubs/runner/metaflow_runner.pyi +3 -3
  161. metaflow-stubs/runner/nbdeploy.pyi +2 -2
  162. metaflow-stubs/runner/nbrun.pyi +2 -2
  163. metaflow-stubs/runner/subprocess_manager.pyi +2 -2
  164. metaflow-stubs/runner/utils.pyi +2 -2
  165. metaflow-stubs/system/__init__.pyi +2 -2
  166. metaflow-stubs/system/system_logger.pyi +2 -2
  167. metaflow-stubs/system/system_monitor.pyi +2 -2
  168. metaflow-stubs/tagging_util.pyi +2 -2
  169. metaflow-stubs/tuple_util.pyi +2 -2
  170. {ob_metaflow_stubs-6.0.3.118.dist-info → ob_metaflow_stubs-6.0.3.120.dist-info}/METADATA +1 -1
  171. ob_metaflow_stubs-6.0.3.120.dist-info/RECORD +174 -0
  172. {ob_metaflow_stubs-6.0.3.118.dist-info → ob_metaflow_stubs-6.0.3.120.dist-info}/WHEEL +1 -1
  173. ob_metaflow_stubs-6.0.3.118.dist-info/RECORD +0 -174
  174. {ob_metaflow_stubs-6.0.3.118.dist-info → ob_metaflow_stubs-6.0.3.120.dist-info}/top_level.txt +0 -0
@@ -1,7 +1,7 @@
1
1
  ######################################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
- # MF version: 2.12.30.1+obcheckpoint(0.1.4);ob(v1) #
4
- # Generated on 2024-11-13T23:34:52.932934 #
3
+ # MF version: 2.12.32.1+obcheckpoint(0.1.4);ob(v1) #
4
+ # Generated on 2024-11-26T19:30:33.284631 #
5
5
  ######################################################################################################
6
6
 
7
7
  from __future__ import annotations
@@ -29,8 +29,8 @@ from .flowspec import FlowSpec as FlowSpec
29
29
  from .parameters import Parameter as Parameter
30
30
  from .parameters import JSONTypeClass as JSONTypeClass
31
31
  from .parameters import JSONType as JSONType
32
- from . import tuple_util as tuple_util
33
32
  from . import cards as cards
33
+ from . import tuple_util as tuple_util
34
34
  from . import events as events
35
35
  from . import runner as runner
36
36
  from . import plugins as plugins
@@ -141,89 +141,121 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
141
141
  ...
142
142
 
143
143
  @typing.overload
144
- def conda(*, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
144
+ 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]]]:
145
145
  """
146
- Specifies the Conda environment for the step.
146
+ Specifies the number of times the task corresponding
147
+ to a step needs to be retried.
147
148
 
148
- Information in this decorator will augment any
149
- attributes set in the `@conda_base` flow-level decorator. Hence,
150
- you can use `@conda_base` to set packages required by all
151
- steps and use `@conda` to specify step-specific overrides.
149
+ This decorator is useful for handling transient errors, such as networking issues.
150
+ If your task contains operations that can't be retried safely, e.g. database updates,
151
+ it is advisable to annotate it with `@retry(times=0)`.
152
+
153
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
154
+ decorator will execute a no-op task after all retries have been exhausted,
155
+ ensuring that the flow execution can continue.
152
156
  """
153
157
  ...
154
158
 
155
159
  @typing.overload
156
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
160
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
157
161
  ...
158
162
 
159
163
  @typing.overload
160
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
164
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
161
165
  ...
162
166
 
163
- def conda(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
167
+ 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):
164
168
  """
165
- Specifies the Conda environment for the step.
169
+ Specifies the number of times the task corresponding
170
+ to a step needs to be retried.
166
171
 
167
- Information in this decorator will augment any
168
- attributes set in the `@conda_base` flow-level decorator. Hence,
169
- you can use `@conda_base` to set packages required by all
170
- steps and use `@conda` to specify step-specific overrides.
172
+ This decorator is useful for handling transient errors, such as networking issues.
173
+ If your task contains operations that can't be retried safely, e.g. database updates,
174
+ it is advisable to annotate it with `@retry(times=0)`.
175
+
176
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
177
+ decorator will execute a no-op task after all retries have been exhausted,
178
+ ensuring that the flow execution can continue.
171
179
  """
172
180
  ...
173
181
 
174
182
  @typing.overload
175
- 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]]]:
183
+ 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]]]:
176
184
  """
177
- Specifies the PyPI packages for the step.
185
+ Specifies the resources needed when executing this step.
178
186
 
179
- Information in this decorator will augment any
180
- attributes set in the `@pyi_base` flow-level decorator. Hence,
181
- you can use `@pypi_base` to set packages required by all
182
- steps and use `@pypi` to specify step-specific overrides.
187
+ Use `@resources` to specify the resource requirements
188
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
189
+
190
+ You can choose the compute layer on the command line by executing e.g.
191
+ ```
192
+ python myflow.py run --with batch
193
+ ```
194
+ or
195
+ ```
196
+ python myflow.py run --with kubernetes
197
+ ```
198
+ which executes the flow on the desired system using the
199
+ requirements specified in `@resources`.
183
200
  """
184
201
  ...
185
202
 
186
203
  @typing.overload
187
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
204
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
188
205
  ...
189
206
 
190
207
  @typing.overload
191
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
208
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
192
209
  ...
193
210
 
194
- 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):
211
+ 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):
195
212
  """
196
- Specifies the PyPI packages for the step.
213
+ Specifies the resources needed when executing this step.
197
214
 
198
- Information in this decorator will augment any
199
- attributes set in the `@pyi_base` flow-level decorator. Hence,
200
- you can use `@pypi_base` to set packages required by all
201
- steps and use `@pypi` to specify step-specific overrides.
215
+ Use `@resources` to specify the resource requirements
216
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
217
+
218
+ You can choose the compute layer on the command line by executing e.g.
219
+ ```
220
+ python myflow.py run --with batch
221
+ ```
222
+ or
223
+ ```
224
+ python myflow.py run --with kubernetes
225
+ ```
226
+ which executes the flow on the desired system using the
227
+ requirements specified in `@resources`.
202
228
  """
203
229
  ...
204
230
 
205
- 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]]]:
231
+ @typing.overload
232
+ def conda(*, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
206
233
  """
207
- Decorator that helps cache, version and store models/datasets from huggingface hub.
234
+ Specifies the Conda environment for the step.
235
+
236
+ Information in this decorator will augment any
237
+ attributes set in the `@conda_base` flow-level decorator. Hence,
238
+ you can use `@conda_base` to set packages required by all
239
+ steps and use `@conda` to specify step-specific overrides.
208
240
  """
209
241
  ...
210
242
 
211
243
  @typing.overload
212
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
213
- """
214
- Decorator prototype for all step decorators. This function gets specialized
215
- and imported for all decorators types by _import_plugin_decorators().
216
- """
244
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
217
245
  ...
218
246
 
219
247
  @typing.overload
220
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
248
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
221
249
  ...
222
250
 
223
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
251
+ def conda(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
224
252
  """
225
- Decorator prototype for all step decorators. This function gets specialized
226
- and imported for all decorators types by _import_plugin_decorators().
253
+ Specifies the Conda environment for the step.
254
+
255
+ Information in this decorator will augment any
256
+ attributes set in the `@conda_base` flow-level decorator. Hence,
257
+ you can use `@conda_base` to set packages required by all
258
+ steps and use `@conda` to specify step-specific overrides.
227
259
  """
228
260
  ...
229
261
 
@@ -259,51 +291,33 @@ def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
259
291
  ...
260
292
 
261
293
  @typing.overload
262
- 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]]]:
294
+ 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]]]:
263
295
  """
264
- Specifies the resources needed when executing this step.
265
-
266
- Use `@resources` to specify the resource requirements
267
- independently of the specific compute layer (`@batch`, `@kubernetes`).
296
+ Specifies the PyPI packages for the step.
268
297
 
269
- You can choose the compute layer on the command line by executing e.g.
270
- ```
271
- python myflow.py run --with batch
272
- ```
273
- or
274
- ```
275
- python myflow.py run --with kubernetes
276
- ```
277
- which executes the flow on the desired system using the
278
- requirements specified in `@resources`.
298
+ Information in this decorator will augment any
299
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
300
+ you can use `@pypi_base` to set packages required by all
301
+ steps and use `@pypi` to specify step-specific overrides.
279
302
  """
280
303
  ...
281
304
 
282
305
  @typing.overload
283
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
306
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
284
307
  ...
285
308
 
286
309
  @typing.overload
287
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
310
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
288
311
  ...
289
312
 
290
- 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):
313
+ 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):
291
314
  """
292
- Specifies the resources needed when executing this step.
293
-
294
- Use `@resources` to specify the resource requirements
295
- independently of the specific compute layer (`@batch`, `@kubernetes`).
315
+ Specifies the PyPI packages for the step.
296
316
 
297
- You can choose the compute layer on the command line by executing e.g.
298
- ```
299
- python myflow.py run --with batch
300
- ```
301
- or
302
- ```
303
- python myflow.py run --with kubernetes
304
- ```
305
- which executes the flow on the desired system using the
306
- requirements specified in `@resources`.
317
+ Information in this decorator will augment any
318
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
319
+ you can use `@pypi_base` to set packages required by all
320
+ steps and use `@pypi` to specify step-specific overrides.
307
321
  """
308
322
  ...
309
323
 
@@ -329,129 +343,140 @@ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], Non
329
343
  ...
330
344
 
331
345
  @typing.overload
332
- 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]]]:
346
+ 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]]]:
333
347
  """
334
- Specifies the number of times the task corresponding
335
- to a step needs to be retried.
336
-
337
- This decorator is useful for handling transient errors, such as networking issues.
338
- If your task contains operations that can't be retried safely, e.g. database updates,
339
- it is advisable to annotate it with `@retry(times=0)`.
340
-
341
- This can be used in conjunction with the `@catch` decorator. The `@catch`
342
- decorator will execute a no-op task after all retries have been exhausted,
343
- ensuring that the flow execution can continue.
348
+ Enables checkpointing for a step.
344
349
  """
345
350
  ...
346
351
 
347
352
  @typing.overload
348
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
353
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
349
354
  ...
350
355
 
351
356
  @typing.overload
352
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
357
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
353
358
  ...
354
359
 
355
- 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):
360
+ 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):
356
361
  """
357
- Specifies the number of times the task corresponding
358
- to a step needs to be retried.
359
-
360
- This decorator is useful for handling transient errors, such as networking issues.
361
- If your task contains operations that can't be retried safely, e.g. database updates,
362
- it is advisable to annotate it with `@retry(times=0)`.
363
-
364
- This can be used in conjunction with the `@catch` decorator. The `@catch`
365
- decorator will execute a no-op task after all retries have been exhausted,
366
- ensuring that the flow execution can continue.
362
+ Enables checkpointing for a step.
367
363
  """
368
364
  ...
369
365
 
370
- @typing.overload
371
- 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]]]:
366
+ 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]]]:
372
367
  """
373
- Specifies secrets to be retrieved and injected as environment variables prior to
374
- the execution of a step.
368
+ Decorator that helps cache, version and store models/datasets from huggingface hub.
375
369
  """
376
370
  ...
377
371
 
378
372
  @typing.overload
379
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
373
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
374
+ """
375
+ Internal decorator to support Fast bakery
376
+ """
380
377
  ...
381
378
 
382
379
  @typing.overload
383
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
380
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
384
381
  ...
385
382
 
386
- 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]]] = []):
383
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
387
384
  """
388
- Specifies secrets to be retrieved and injected as environment variables prior to
389
- the execution of a step.
385
+ Internal decorator to support Fast bakery
386
+ """
387
+ ...
388
+
389
+ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: typing.Optional[str] = None, image_pull_policy: str = 'KUBERNETES_IMAGE_PULL_POLICY', service_account: str = 'METAFLOW_KUBERNETES_SERVICE_ACCOUNT', secrets: typing.Optional[typing.List[str]] = None, node_selector: typing.Union[typing.Dict[str, str], str, None] = None, namespace: str = 'METAFLOW_KUBERNETES_NAMESPACE', gpu: typing.Optional[int] = None, gpu_vendor: str = 'KUBERNETES_GPU_VENDOR', tolerations: typing.List[str] = [], use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = '/metaflow_temp', persistent_volume_claims: typing.Optional[typing.Dict[str, str]] = None, shared_memory: typing.Optional[int] = None, port: typing.Optional[int] = None, compute_pool: typing.Optional[str] = None, hostname_resolution_timeout: int = 600) -> 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]]]:
390
+ """
391
+ Specifies that this step should execute on Kubernetes.
390
392
  """
391
393
  ...
392
394
 
393
395
  @typing.overload
394
- 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]]]:
396
+ 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]]]:
395
397
  """
396
- Enables loading / saving of models within a step.
398
+ Creates a human-readable report, a Metaflow Card, after this step completes.
399
+
400
+ Note that you may add multiple `@card` decorators in a step with different parameters.
397
401
  """
398
402
  ...
399
403
 
400
404
  @typing.overload
401
- def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
405
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
402
406
  ...
403
407
 
404
408
  @typing.overload
405
- def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
409
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
406
410
  ...
407
411
 
408
- 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):
412
+ 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):
409
413
  """
410
- Enables loading / saving of models within a step.
414
+ Creates a human-readable report, a Metaflow Card, after this step completes.
415
+
416
+ Note that you may add multiple `@card` decorators in a step with different parameters.
411
417
  """
412
418
  ...
413
419
 
414
420
  @typing.overload
415
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
421
+ 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]]]:
416
422
  """
417
- Internal decorator to support Fast bakery
423
+ Specifies secrets to be retrieved and injected as environment variables prior to
424
+ the execution of a step.
418
425
  """
419
426
  ...
420
427
 
421
428
  @typing.overload
422
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
429
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
423
430
  ...
424
431
 
425
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
432
+ @typing.overload
433
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
434
+ ...
435
+
436
+ 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]]] = []):
426
437
  """
427
- Internal decorator to support Fast bakery
438
+ Specifies secrets to be retrieved and injected as environment variables prior to
439
+ the execution of a step.
428
440
  """
429
441
  ...
430
442
 
431
- def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: typing.Optional[str] = None, image_pull_policy: str = 'KUBERNETES_IMAGE_PULL_POLICY', service_account: str = 'METAFLOW_KUBERNETES_SERVICE_ACCOUNT', secrets: typing.Optional[typing.List[str]] = None, node_selector: typing.Union[typing.Dict[str, str], str, None] = None, namespace: str = 'METAFLOW_KUBERNETES_NAMESPACE', gpu: typing.Optional[int] = None, gpu_vendor: str = 'KUBERNETES_GPU_VENDOR', tolerations: typing.List[str] = [], use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = '/metaflow_temp', persistent_volume_claims: typing.Optional[typing.Dict[str, str]] = None, shared_memory: typing.Optional[int] = None, port: typing.Optional[int] = None, compute_pool: typing.Optional[str] = None, hostname_resolution_timeout: int = 600) -> 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]]]:
443
+ @typing.overload
444
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
432
445
  """
433
- Specifies that this step should execute on Kubernetes.
446
+ Decorator prototype for all step decorators. This function gets specialized
447
+ and imported for all decorators types by _import_plugin_decorators().
434
448
  """
435
449
  ...
436
450
 
437
451
  @typing.overload
438
- 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]]]:
452
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
453
+ ...
454
+
455
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
439
456
  """
440
- Enables checkpointing for a step.
457
+ Decorator prototype for all step decorators. This function gets specialized
458
+ and imported for all decorators types by _import_plugin_decorators().
441
459
  """
442
460
  ...
443
461
 
444
462
  @typing.overload
445
- def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
463
+ 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]]]:
464
+ """
465
+ Enables loading / saving of models within a step.
466
+ """
446
467
  ...
447
468
 
448
469
  @typing.overload
449
- def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
470
+ def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
450
471
  ...
451
472
 
452
- 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):
473
+ @typing.overload
474
+ def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
475
+ ...
476
+
477
+ 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):
453
478
  """
454
- Enables checkpointing for a step.
479
+ Enables loading / saving of models within a step.
455
480
  """
456
481
  ...
457
482
 
@@ -494,78 +519,6 @@ def timeout(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None],
494
519
  """
495
520
  ...
496
521
 
497
- @typing.overload
498
- 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]]]:
499
- """
500
- Creates a human-readable report, a Metaflow Card, after this step completes.
501
-
502
- Note that you may add multiple `@card` decorators in a step with different parameters.
503
- """
504
- ...
505
-
506
- @typing.overload
507
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
508
- ...
509
-
510
- @typing.overload
511
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
512
- ...
513
-
514
- 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):
515
- """
516
- Creates a human-readable report, a Metaflow Card, after this step completes.
517
-
518
- Note that you may add multiple `@card` decorators in a step with different parameters.
519
- """
520
- ...
521
-
522
- def nim(*, models: "list[NIM]", backend: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
523
- """
524
- This decorator is used to run NIM containers in Metaflow tasks as sidecars.
525
-
526
- User code call
527
- -----------
528
- @nim(
529
- models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
530
- backend='managed'
531
- )
532
-
533
- Valid backend options
534
- ---------------------
535
- - 'managed': Outerbounds selects a compute provider based on the model.
536
- - 🚧 'dataplane': Run in your account.
537
-
538
- Valid model options
539
- ----------------
540
- - 'meta/llama3-8b-instruct': 8B parameter model
541
- - 'meta/llama3-70b-instruct': 70B parameter model
542
- - Upon request, any model here: https://nvcf.ngc.nvidia.com/functions?filter=nvidia-functions
543
- """
544
- ...
545
-
546
- @typing.overload
547
- 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]]:
548
- """
549
- Specifies the Conda environment for all steps of the flow.
550
-
551
- Use `@conda_base` to set common libraries required by all
552
- steps and use `@conda` to specify step-specific additions.
553
- """
554
- ...
555
-
556
- @typing.overload
557
- def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
558
- ...
559
-
560
- 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):
561
- """
562
- Specifies the Conda environment for all steps of the flow.
563
-
564
- Use `@conda_base` to set common libraries required by all
565
- steps and use `@conda` to specify step-specific additions.
566
- """
567
- ...
568
-
569
522
  @typing.overload
570
523
  def schedule(*, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
571
524
  """
@@ -585,45 +538,6 @@ def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly:
585
538
  """
586
539
  ...
587
540
 
588
- def project(*, name: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
589
- """
590
- Specifies what flows belong to the same project.
591
-
592
- A project-specific namespace is created for all flows that
593
- use the same `@project(name)`.
594
- """
595
- ...
596
-
597
- def airflow_external_task_sensor(*, timeout: int, poke_interval: int, mode: str, exponential_backoff: bool, pool: str, soft_fail: bool, name: str, description: str, external_dag_id: str, external_task_ids: typing.List[str], allowed_states: typing.List[str], failed_states: typing.List[str], execution_delta: "datetime.timedelta", check_existence: bool) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
598
- """
599
- The `@airflow_external_task_sensor` decorator attaches a Airflow [ExternalTaskSensor](https://airflow.apache.org/docs/apache-airflow/stable/_api/airflow/sensors/external_task/index.html#airflow.sensors.external_task.ExternalTaskSensor) before the start step of the flow.
600
- This decorator only works when a flow is scheduled on Airflow and is compiled using `airflow create`. More than one `@airflow_external_task_sensor` can be added as a flow decorators. Adding more than one decorator will ensure that `start` step starts only after all sensors finish.
601
- """
602
- ...
603
-
604
- @typing.overload
605
- def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
606
- """
607
- Specifies the PyPI packages for all steps of the flow.
608
-
609
- Use `@pypi_base` to set common packages required by all
610
- steps and use `@pypi` to specify step-specific overrides.
611
- """
612
- ...
613
-
614
- @typing.overload
615
- def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
616
- ...
617
-
618
- def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
619
- """
620
- Specifies the PyPI packages for all steps of the flow.
621
-
622
- Use `@pypi_base` to set common packages required by all
623
- steps and use `@pypi` to specify step-specific overrides.
624
- """
625
- ...
626
-
627
541
  @typing.overload
628
542
  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]]:
629
543
  """
@@ -705,13 +619,10 @@ def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *
705
619
  """
706
620
  ...
707
621
 
708
- 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]]:
622
+ def airflow_external_task_sensor(*, timeout: int, poke_interval: int, mode: str, exponential_backoff: bool, pool: str, soft_fail: bool, name: str, description: str, external_dag_id: str, external_task_ids: typing.List[str], allowed_states: typing.List[str], failed_states: typing.List[str], execution_delta: "datetime.timedelta", check_existence: bool) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
709
623
  """
710
- 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)
711
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
712
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
713
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
714
- starts only after all sensors finish.
624
+ The `@airflow_external_task_sensor` decorator attaches a Airflow [ExternalTaskSensor](https://airflow.apache.org/docs/apache-airflow/stable/_api/airflow/sensors/external_task/index.html#airflow.sensors.external_task.ExternalTaskSensor) before the start step of the flow.
625
+ This decorator only works when a flow is scheduled on Airflow and is compiled using `airflow create`. More than one `@airflow_external_task_sensor` can be added as a flow decorators. Adding more than one decorator will ensure that `start` step starts only after all sensors finish.
715
626
  """
716
627
  ...
717
628
 
@@ -788,5 +699,94 @@ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: t
788
699
  """
789
700
  ...
790
701
 
702
+ @typing.overload
703
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
704
+ """
705
+ Specifies the PyPI packages for all steps of the flow.
706
+
707
+ Use `@pypi_base` to set common packages required by all
708
+ steps and use `@pypi` to specify step-specific overrides.
709
+ """
710
+ ...
711
+
712
+ @typing.overload
713
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
714
+ ...
715
+
716
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
717
+ """
718
+ Specifies the PyPI packages for all steps of the flow.
719
+
720
+ Use `@pypi_base` to set common packages required by all
721
+ steps and use `@pypi` to specify step-specific overrides.
722
+ """
723
+ ...
724
+
725
+ @typing.overload
726
+ 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]]:
727
+ """
728
+ Specifies the Conda environment for all steps of the flow.
729
+
730
+ Use `@conda_base` to set common libraries required by all
731
+ steps and use `@conda` to specify step-specific additions.
732
+ """
733
+ ...
734
+
735
+ @typing.overload
736
+ def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
737
+ ...
738
+
739
+ 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):
740
+ """
741
+ Specifies the Conda environment for all steps of the flow.
742
+
743
+ Use `@conda_base` to set common libraries required by all
744
+ steps and use `@conda` to specify step-specific additions.
745
+ """
746
+ ...
747
+
748
+ 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]]:
749
+ """
750
+ 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)
751
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
752
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
753
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
754
+ starts only after all sensors finish.
755
+ """
756
+ ...
757
+
758
+ def nim(*, models: "list[NIM]", backend: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
759
+ """
760
+ This decorator is used to run NIM containers in Metaflow tasks as sidecars.
761
+
762
+ User code call
763
+ -----------
764
+ @nim(
765
+ models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
766
+ backend='managed'
767
+ )
768
+
769
+ Valid backend options
770
+ ---------------------
771
+ - 'managed': Outerbounds selects a compute provider based on the model.
772
+ - 🚧 'dataplane': Run in your account.
773
+
774
+ Valid model options
775
+ ----------------
776
+ - 'meta/llama3-8b-instruct': 8B parameter model
777
+ - 'meta/llama3-70b-instruct': 70B parameter model
778
+ - Upon request, any model here: https://nvcf.ngc.nvidia.com/functions?filter=nvidia-functions
779
+ """
780
+ ...
781
+
782
+ def project(*, name: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
783
+ """
784
+ Specifies what flows belong to the same project.
785
+
786
+ A project-specific namespace is created for all flows that
787
+ use the same `@project(name)`.
788
+ """
789
+ ...
790
+
791
791
  pkg_name: str
792
792