ob-metaflow-stubs 6.0.3.156__py2.py3-none-any.whl → 6.0.3.158__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 (200) hide show
  1. metaflow-stubs/__init__.pyi +741 -626
  2. metaflow-stubs/cards.pyi +2 -2
  3. metaflow-stubs/cli.pyi +2 -2
  4. metaflow-stubs/cli_components/__init__.pyi +2 -2
  5. metaflow-stubs/cli_components/utils.pyi +2 -2
  6. metaflow-stubs/client/__init__.pyi +2 -2
  7. metaflow-stubs/client/core.pyi +6 -6
  8. metaflow-stubs/client/filecache.pyi +3 -3
  9. metaflow-stubs/events.pyi +3 -3
  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 +3 -3
  18. metaflow-stubs/metadata_provider/util.pyi +2 -2
  19. metaflow-stubs/metaflow_config.pyi +2 -2
  20. metaflow-stubs/metaflow_current.pyi +37 -37
  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 +4 -4
  32. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/lineage_card.pyi +2 -2
  33. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/checkpoint_storage.pyi +6 -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 +11 -5
  36. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/decorator.pyi +5 -4
  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 +5 -3
  39. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/lineage.pyi +2 -2
  40. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/__init__.pyi +2 -2
  41. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/context.pyi +95 -0
  42. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/core.pyi +7 -9
  43. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/decorator.pyi +135 -0
  44. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/exceptions.pyi +2 -2
  45. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/task_utils.pyi +33 -5
  46. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/utils.pyi +2 -2
  47. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastructures.pyi +3 -3
  48. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/exceptions.pyi +2 -2
  49. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/__init__.pyi +2 -2
  50. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/decorator.pyi +3 -3
  51. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/__init__.pyi +2 -2
  52. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/core.pyi +2 -2
  53. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/exceptions.pyi +2 -2
  54. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/model_storage.pyi +3 -3
  55. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/__init__.pyi +2 -2
  56. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/flowspec_utils.pyi +2 -2
  57. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/general.pyi +2 -2
  58. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/identity_utils.pyi +3 -3
  59. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/__init__.pyi +2 -2
  60. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/base.pyi +2 -2
  61. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/tar.pyi +3 -3
  62. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/tar_utils.pyi +3 -3
  63. metaflow-stubs/mf_extensions/outerbounds/__init__.pyi +2 -2
  64. metaflow-stubs/mf_extensions/outerbounds/plugins/__init__.pyi +2 -2
  65. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/__init__.pyi +2 -2
  66. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/injector.pyi +2 -2
  67. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/__init__.pyi +2 -2
  68. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/ollama.pyi +2 -2
  69. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/__init__.pyi +2 -2
  70. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/snowflake.pyi +2 -2
  71. metaflow-stubs/mf_extensions/outerbounds/profilers/__init__.pyi +2 -2
  72. metaflow-stubs/mf_extensions/outerbounds/profilers/gpu.pyi +2 -2
  73. metaflow-stubs/mf_extensions/outerbounds/remote_config.pyi +2 -2
  74. metaflow-stubs/mf_extensions/outerbounds/toplevel/__init__.pyi +2 -2
  75. metaflow-stubs/mf_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.pyi +2 -2
  76. metaflow-stubs/multicore_utils.pyi +2 -2
  77. metaflow-stubs/parameters.pyi +4 -4
  78. metaflow-stubs/plugins/__init__.pyi +11 -11
  79. metaflow-stubs/plugins/airflow/__init__.pyi +2 -2
  80. metaflow-stubs/plugins/airflow/airflow_utils.pyi +2 -2
  81. metaflow-stubs/plugins/airflow/exception.pyi +2 -2
  82. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +2 -2
  83. metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi +2 -2
  84. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +2 -2
  85. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +2 -2
  86. metaflow-stubs/plugins/argo/__init__.pyi +2 -2
  87. metaflow-stubs/plugins/argo/argo_client.pyi +2 -2
  88. metaflow-stubs/plugins/argo/argo_events.pyi +2 -2
  89. metaflow-stubs/plugins/argo/argo_workflows.pyi +4 -4
  90. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +2 -2
  91. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +5 -5
  92. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +3 -3
  93. metaflow-stubs/plugins/aws/__init__.pyi +2 -2
  94. metaflow-stubs/plugins/aws/aws_client.pyi +2 -2
  95. metaflow-stubs/plugins/aws/aws_utils.pyi +2 -2
  96. metaflow-stubs/plugins/aws/batch/__init__.pyi +2 -2
  97. metaflow-stubs/plugins/aws/batch/batch.pyi +2 -2
  98. metaflow-stubs/plugins/aws/batch/batch_client.pyi +2 -2
  99. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +2 -2
  100. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +2 -2
  101. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +5 -5
  102. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +2 -2
  103. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +2 -2
  104. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +2 -2
  105. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +2 -2
  106. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +2 -2
  107. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer.pyi +4 -4
  108. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer_objects.pyi +3 -3
  109. metaflow-stubs/plugins/azure/__init__.pyi +2 -2
  110. metaflow-stubs/plugins/azure/azure_credential.pyi +2 -2
  111. metaflow-stubs/plugins/azure/azure_exceptions.pyi +2 -2
  112. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +5 -5
  113. metaflow-stubs/plugins/azure/azure_utils.pyi +2 -2
  114. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +2 -2
  115. metaflow-stubs/plugins/azure/includefile_support.pyi +2 -2
  116. metaflow-stubs/plugins/cards/__init__.pyi +6 -6
  117. metaflow-stubs/plugins/cards/card_client.pyi +3 -3
  118. metaflow-stubs/plugins/cards/card_creator.pyi +2 -2
  119. metaflow-stubs/plugins/cards/card_datastore.pyi +2 -2
  120. metaflow-stubs/plugins/cards/card_decorator.pyi +2 -2
  121. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +2 -2
  122. metaflow-stubs/plugins/cards/card_modules/basic.pyi +3 -3
  123. metaflow-stubs/plugins/cards/card_modules/card.pyi +2 -2
  124. metaflow-stubs/plugins/cards/card_modules/components.pyi +5 -5
  125. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +2 -2
  126. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +2 -2
  127. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +2 -2
  128. metaflow-stubs/plugins/cards/card_resolver.pyi +2 -2
  129. metaflow-stubs/plugins/cards/component_serializer.pyi +2 -2
  130. metaflow-stubs/plugins/cards/exception.pyi +2 -2
  131. metaflow-stubs/plugins/catch_decorator.pyi +2 -2
  132. metaflow-stubs/plugins/datatools/__init__.pyi +2 -2
  133. metaflow-stubs/plugins/datatools/local.pyi +2 -2
  134. metaflow-stubs/plugins/datatools/s3/__init__.pyi +2 -2
  135. metaflow-stubs/plugins/datatools/s3/s3.pyi +4 -4
  136. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +2 -2
  137. metaflow-stubs/plugins/datatools/s3/s3util.pyi +2 -2
  138. metaflow-stubs/plugins/debug_logger.pyi +2 -2
  139. metaflow-stubs/plugins/debug_monitor.pyi +2 -2
  140. metaflow-stubs/plugins/environment_decorator.pyi +2 -2
  141. metaflow-stubs/plugins/events_decorator.pyi +2 -2
  142. metaflow-stubs/plugins/frameworks/__init__.pyi +2 -2
  143. metaflow-stubs/plugins/frameworks/pytorch.pyi +2 -2
  144. metaflow-stubs/plugins/gcp/__init__.pyi +2 -2
  145. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +5 -5
  146. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +2 -2
  147. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +2 -2
  148. metaflow-stubs/plugins/gcp/gs_utils.pyi +2 -2
  149. metaflow-stubs/plugins/gcp/includefile_support.pyi +2 -2
  150. metaflow-stubs/plugins/kubernetes/__init__.pyi +2 -2
  151. metaflow-stubs/plugins/kubernetes/kube_utils.pyi +3 -3
  152. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +2 -2
  153. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +2 -2
  154. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +2 -2
  155. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +2 -2
  156. metaflow-stubs/plugins/kubernetes/spot_monitor_sidecar.pyi +2 -2
  157. metaflow-stubs/plugins/ollama/__init__.pyi +2 -2
  158. metaflow-stubs/plugins/parallel_decorator.pyi +2 -2
  159. metaflow-stubs/plugins/perimeters.pyi +2 -2
  160. metaflow-stubs/plugins/project_decorator.pyi +2 -2
  161. metaflow-stubs/plugins/pypi/__init__.pyi +3 -3
  162. metaflow-stubs/plugins/pypi/conda_decorator.pyi +2 -2
  163. metaflow-stubs/plugins/pypi/conda_environment.pyi +5 -5
  164. metaflow-stubs/plugins/pypi/parsers.pyi +2 -2
  165. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +2 -2
  166. metaflow-stubs/plugins/pypi/pypi_environment.pyi +2 -2
  167. metaflow-stubs/plugins/pypi/utils.pyi +2 -2
  168. metaflow-stubs/plugins/resources_decorator.pyi +2 -2
  169. metaflow-stubs/plugins/retry_decorator.pyi +2 -2
  170. metaflow-stubs/plugins/secrets/__init__.pyi +2 -2
  171. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +4 -4
  172. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +2 -2
  173. metaflow-stubs/plugins/snowflake/__init__.pyi +2 -2
  174. metaflow-stubs/plugins/storage_executor.pyi +2 -2
  175. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +2 -2
  176. metaflow-stubs/plugins/timeout_decorator.pyi +2 -2
  177. metaflow-stubs/profilers/__init__.pyi +2 -2
  178. metaflow-stubs/pylint_wrapper.pyi +2 -2
  179. metaflow-stubs/runner/__init__.pyi +2 -2
  180. metaflow-stubs/runner/deployer.pyi +31 -31
  181. metaflow-stubs/runner/deployer_impl.pyi +3 -3
  182. metaflow-stubs/runner/metaflow_runner.pyi +4 -4
  183. metaflow-stubs/runner/nbdeploy.pyi +2 -2
  184. metaflow-stubs/runner/nbrun.pyi +2 -2
  185. metaflow-stubs/runner/subprocess_manager.pyi +2 -2
  186. metaflow-stubs/runner/utils.pyi +3 -3
  187. metaflow-stubs/system/__init__.pyi +2 -2
  188. metaflow-stubs/system/system_logger.pyi +3 -3
  189. metaflow-stubs/system/system_monitor.pyi +2 -2
  190. metaflow-stubs/tagging_util.pyi +2 -2
  191. metaflow-stubs/tuple_util.pyi +2 -2
  192. metaflow-stubs/user_configs/__init__.pyi +2 -2
  193. metaflow-stubs/user_configs/config_decorators.pyi +7 -7
  194. metaflow-stubs/user_configs/config_options.pyi +4 -4
  195. metaflow-stubs/user_configs/config_parameters.pyi +5 -5
  196. {ob_metaflow_stubs-6.0.3.156.dist-info → ob_metaflow_stubs-6.0.3.158.dist-info}/METADATA +1 -1
  197. ob_metaflow_stubs-6.0.3.158.dist-info/RECORD +200 -0
  198. ob_metaflow_stubs-6.0.3.156.dist-info/RECORD +0 -198
  199. {ob_metaflow_stubs-6.0.3.156.dist-info → ob_metaflow_stubs-6.0.3.158.dist-info}/WHEEL +0 -0
  200. {ob_metaflow_stubs-6.0.3.156.dist-info → ob_metaflow_stubs-6.0.3.158.dist-info}/top_level.txt +0 -0
@@ -1,15 +1,15 @@
1
1
  ######################################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
- # MF version: 2.15.7.2+obcheckpoint(0.1.9);ob(v1) #
4
- # Generated on 2025-04-16T21:50:24.175279 #
3
+ # MF version: 2.15.7.2+obcheckpoint(0.2.1);ob(v1) #
4
+ # Generated on 2025-04-17T20:45:30.855915 #
5
5
  ######################################################################################################
6
6
 
7
7
  from __future__ import annotations
8
8
 
9
9
  import typing
10
10
  if typing.TYPE_CHECKING:
11
- import datetime
12
11
  import typing
12
+ import datetime
13
13
  FlowSpecDerived = typing.TypeVar("FlowSpecDerived", bound="FlowSpec", contravariant=False, covariant=False)
14
14
  StepFlag = typing.NewType("StepFlag", bool)
15
15
 
@@ -35,8 +35,8 @@ from .user_configs.config_parameters import ConfigValue as ConfigValue
35
35
  from .user_configs.config_parameters import config_expr as config_expr
36
36
  from .user_configs.config_decorators import CustomFlowDecorator as CustomFlowDecorator
37
37
  from .user_configs.config_decorators import CustomStepDecorator as CustomStepDecorator
38
- from . import tuple_util as tuple_util
39
38
  from . import cards as cards
39
+ from . import tuple_util as tuple_util
40
40
  from . import events as events
41
41
  from . import runner as runner
42
42
  from . import plugins as plugins
@@ -44,8 +44,8 @@ from .mf_extensions.outerbounds.toplevel.global_aliases_for_metaflow_package imp
44
44
  from . import includefile as includefile
45
45
  from .includefile import IncludeFile as IncludeFile
46
46
  from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
47
- from .plugins.pypi.parsers import conda_environment_yml_parser as conda_environment_yml_parser
48
47
  from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
48
+ from .plugins.pypi.parsers import conda_environment_yml_parser as conda_environment_yml_parser
49
49
  from . import client as client
50
50
  from .client.core import namespace as namespace
51
51
  from .client.core import get_namespace as get_namespace
@@ -66,6 +66,7 @@ from .runner.deployer import DeployedFlow as DeployedFlow
66
66
  from .runner.nbdeploy import NBDeployer as NBDeployer
67
67
  from .mf_extensions.obcheckpoint.plugins.machine_learning_utilities.checkpoints.final_api import Checkpoint as Checkpoint
68
68
  from .mf_extensions.obcheckpoint.plugins.machine_learning_utilities.datastructures import load_model as load_model
69
+ from .mf_extensions.obcheckpoint.plugins.machine_learning_utilities.datastore.context import artifact_store_from as artifact_store_from
69
70
  from .mf_extensions.outerbounds.toplevel.global_aliases_for_metaflow_package import get_aws_client as get_aws_client
70
71
  from .mf_extensions.outerbounds.plugins.snowflake.snowflake import Snowflake as Snowflake
71
72
  from . import cli_components as cli_components
@@ -149,22 +150,6 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
149
150
  """
150
151
  ...
151
152
 
152
- 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]]]:
153
- """
154
- Specifies that this step should execute on DGX cloud.
155
-
156
-
157
- Parameters
158
- ----------
159
- gpu : int
160
- Number of GPUs to use.
161
- gpu_type : str
162
- Type of Nvidia GPU to use.
163
- queue_timeout : int
164
- Time to keep the job in NVCF's queue.
165
- """
166
- ...
167
-
168
153
  @typing.overload
169
154
  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]]]:
170
155
  """
@@ -217,81 +202,194 @@ def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
217
202
  ...
218
203
 
219
204
  @typing.overload
220
- 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]]]:
205
+ 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]]]:
221
206
  """
222
- Specifies the resources needed when executing this step.
223
-
224
- Use `@resources` to specify the resource requirements
225
- independently of the specific compute layer (`@batch`, `@kubernetes`).
207
+ Creates a human-readable report, a Metaflow Card, after this step completes.
226
208
 
227
- You can choose the compute layer on the command line by executing e.g.
228
- ```
229
- python myflow.py run --with batch
230
- ```
231
- or
232
- ```
233
- python myflow.py run --with kubernetes
234
- ```
235
- which executes the flow on the desired system using the
236
- requirements specified in `@resources`.
209
+ Note that you may add multiple `@card` decorators in a step with different parameters.
237
210
 
238
211
 
239
212
  Parameters
240
213
  ----------
241
- cpu : int, default 1
242
- Number of CPUs required for this step.
243
- gpu : int, optional, default None
244
- Number of GPUs required for this step.
245
- disk : int, optional, default None
246
- Disk size (in MB) required for this step. Only applies on Kubernetes.
247
- memory : int, default 4096
248
- Memory size (in MB) required for this step.
249
- shared_memory : int, optional, default None
250
- The value for the size (in MiB) of the /dev/shm volume for this step.
251
- This parameter maps to the `--shm-size` option in Docker.
214
+ type : str, default 'default'
215
+ Card type.
216
+ id : str, optional, default None
217
+ If multiple cards are present, use this id to identify this card.
218
+ options : Dict[str, Any], default {}
219
+ Options passed to the card. The contents depend on the card type.
220
+ timeout : int, default 45
221
+ Interrupt reporting if it takes more than this many seconds.
252
222
  """
253
223
  ...
254
224
 
255
225
  @typing.overload
256
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
226
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
257
227
  ...
258
228
 
259
229
  @typing.overload
260
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
230
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
261
231
  ...
262
232
 
263
- 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):
233
+ 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):
264
234
  """
265
- Specifies the resources needed when executing this step.
235
+ Creates a human-readable report, a Metaflow Card, after this step completes.
266
236
 
267
- Use `@resources` to specify the resource requirements
268
- independently of the specific compute layer (`@batch`, `@kubernetes`).
237
+ Note that you may add multiple `@card` decorators in a step with different parameters.
269
238
 
270
- You can choose the compute layer on the command line by executing e.g.
271
- ```
272
- python myflow.py run --with batch
273
- ```
274
- or
275
- ```
276
- python myflow.py run --with kubernetes
277
- ```
278
- which executes the flow on the desired system using the
279
- requirements specified in `@resources`.
239
+
240
+ Parameters
241
+ ----------
242
+ type : str, default 'default'
243
+ Card type.
244
+ id : str, optional, default None
245
+ If multiple cards are present, use this id to identify this card.
246
+ options : Dict[str, Any], default {}
247
+ Options passed to the card. The contents depend on the card type.
248
+ timeout : int, default 45
249
+ Interrupt reporting if it takes more than this many seconds.
250
+ """
251
+ ...
252
+
253
+ 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]]]:
254
+ """
255
+ This decorator is used to run Ollama APIs as Metaflow task sidecars.
256
+
257
+ User code call
258
+ -----------
259
+ @ollama(
260
+ models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
261
+ backend='local'
262
+ )
263
+
264
+ Valid backend options
265
+ ---------------------
266
+ - 'local': Run as a separate process on the local task machine.
267
+ - (TODO) 'managed': Outerbounds hosts and selects compute provider.
268
+ - (TODO) 'remote': Spin up separate instance to serve Ollama models.
269
+
270
+ Valid model options
271
+ ----------------
272
+ - 'llama3.2'
273
+ - 'llama3.3'
274
+ - any model here https://ollama.com/search
275
+
276
+
277
+ Parameters
278
+ ----------
279
+ models: list[Ollama]
280
+ List of Ollama containers running models in sidecars.
281
+ backend: str
282
+ Determines where and how to run the Ollama process.
283
+ """
284
+ ...
285
+
286
+ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: typing.Optional[str] = None, image_pull_policy: str = 'KUBERNETES_IMAGE_PULL_POLICY', service_account: str = 'METAFLOW_KUBERNETES_SERVICE_ACCOUNT', secrets: typing.Optional[typing.List[str]] = None, node_selector: typing.Union[typing.Dict[str, str], str, None] = None, namespace: str = 'METAFLOW_KUBERNETES_NAMESPACE', gpu: typing.Optional[int] = None, gpu_vendor: str = 'KUBERNETES_GPU_VENDOR', tolerations: typing.List[str] = [], labels: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_LABELS', annotations: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_ANNOTATIONS', use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = '/metaflow_temp', persistent_volume_claims: typing.Optional[typing.Dict[str, str]] = None, shared_memory: typing.Optional[int] = None, port: typing.Optional[int] = None, compute_pool: typing.Optional[str] = None, hostname_resolution_timeout: int = 600, qos: str = 'Burstable') -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
287
+ """
288
+ Specifies that this step should execute on Kubernetes.
280
289
 
281
290
 
282
291
  Parameters
283
292
  ----------
284
293
  cpu : int, default 1
285
- Number of CPUs required for this step.
286
- gpu : int, optional, default None
287
- Number of GPUs required for this step.
288
- disk : int, optional, default None
289
- Disk size (in MB) required for this step. Only applies on Kubernetes.
294
+ Number of CPUs required for this step. If `@resources` is
295
+ also present, the maximum value from all decorators is used.
290
296
  memory : int, default 4096
291
- Memory size (in MB) required for this step.
292
- shared_memory : int, optional, default None
293
- The value for the size (in MiB) of the /dev/shm volume for this step.
294
- This parameter maps to the `--shm-size` option in Docker.
297
+ Memory size (in MB) required for this step. If
298
+ `@resources` is also present, the maximum value from all decorators is
299
+ used.
300
+ disk : int, default 10240
301
+ Disk size (in MB) required for this step. If
302
+ `@resources` is also present, the maximum value from all decorators is
303
+ used.
304
+ image : str, optional, default None
305
+ Docker image to use when launching on Kubernetes. If not specified, and
306
+ METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
307
+ not, a default Docker image mapping to the current version of Python is used.
308
+ image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
309
+ If given, the imagePullPolicy to be applied to the Docker image of the step.
310
+ service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
311
+ Kubernetes service account to use when launching pod in Kubernetes.
312
+ secrets : List[str], optional, default None
313
+ Kubernetes secrets to use when launching pod in Kubernetes. These
314
+ secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
315
+ in Metaflow configuration.
316
+ node_selector: Union[Dict[str,str], str], optional, default None
317
+ Kubernetes node selector(s) to apply to the pod running the task.
318
+ Can be passed in as a comma separated string of values e.g.
319
+ 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
320
+ {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
321
+ namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
322
+ Kubernetes namespace to use when launching pod in Kubernetes.
323
+ gpu : int, optional, default None
324
+ Number of GPUs required for this step. A value of zero implies that
325
+ the scheduled node should not have GPUs.
326
+ gpu_vendor : str, default KUBERNETES_GPU_VENDOR
327
+ The vendor of the GPUs to be used for this step.
328
+ tolerations : List[str], default []
329
+ The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
330
+ Kubernetes tolerations to use when launching pod in Kubernetes.
331
+ labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
332
+ Kubernetes labels to use when launching pod in Kubernetes.
333
+ annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
334
+ Kubernetes annotations to use when launching pod in Kubernetes.
335
+ use_tmpfs : bool, default False
336
+ This enables an explicit tmpfs mount for this step.
337
+ tmpfs_tempdir : bool, default True
338
+ sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
339
+ tmpfs_size : int, optional, default: None
340
+ The value for the size (in MiB) of the tmpfs mount for this step.
341
+ This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
342
+ memory allocated for this step.
343
+ tmpfs_path : str, optional, default /metaflow_temp
344
+ Path to tmpfs mount for this step.
345
+ persistent_volume_claims : Dict[str, str], optional, default None
346
+ A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
347
+ volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
348
+ shared_memory: int, optional
349
+ Shared memory size (in MiB) required for this step
350
+ port: int, optional
351
+ Port number to specify in the Kubernetes job object
352
+ compute_pool : str, optional, default None
353
+ Compute pool to be used for for this step.
354
+ If not specified, any accessible compute pool within the perimeter is used.
355
+ hostname_resolution_timeout: int, default 10 * 60
356
+ Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
357
+ Only applicable when @parallel is used.
358
+ qos: str, default: Burstable
359
+ Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
360
+ """
361
+ ...
362
+
363
+ def nim(*, models: "list[NIM]", backend: 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]]]:
364
+ """
365
+ This decorator is used to run NIM containers in Metaflow tasks as sidecars.
366
+
367
+ User code call
368
+ -----------
369
+ @nim(
370
+ models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
371
+ backend='managed'
372
+ )
373
+
374
+ Valid backend options
375
+ ---------------------
376
+ - 'managed': Outerbounds selects a compute provider based on the model.
377
+
378
+ Valid model options
379
+ ----------------
380
+ - 'meta/llama3-8b-instruct': 8B parameter model
381
+ - 'meta/llama3-70b-instruct': 70B parameter model
382
+ - any model here: https://nvcf.ngc.nvidia.com/functions?filter=nvidia-functions
383
+
384
+
385
+ Parameters
386
+ ----------
387
+ models: list[NIM]
388
+ List of NIM containers running models in sidecars.
389
+ backend: str
390
+ Compute provider to run the NIM container.
391
+ queue_timeout : int
392
+ Time to keep the job in NVCF's queue.
295
393
  """
296
394
  ...
297
395
 
@@ -353,66 +451,7 @@ def checkpoint(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None
353
451
  ...
354
452
 
355
453
  @typing.overload
356
- 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]]]:
357
- """
358
- Specifies a timeout for your step.
359
-
360
- This decorator is useful if this step may hang indefinitely.
361
-
362
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
363
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
364
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
365
-
366
- Note that all the values specified in parameters are added together so if you specify
367
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
368
-
369
-
370
- Parameters
371
- ----------
372
- seconds : int, default 0
373
- Number of seconds to wait prior to timing out.
374
- minutes : int, default 0
375
- Number of minutes to wait prior to timing out.
376
- hours : int, default 0
377
- Number of hours to wait prior to timing out.
378
- """
379
- ...
380
-
381
- @typing.overload
382
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
383
- ...
384
-
385
- @typing.overload
386
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
387
- ...
388
-
389
- 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):
390
- """
391
- Specifies a timeout for your step.
392
-
393
- This decorator is useful if this step may hang indefinitely.
394
-
395
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
396
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
397
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
398
-
399
- Note that all the values specified in parameters are added together so if you specify
400
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
401
-
402
-
403
- Parameters
404
- ----------
405
- seconds : int, default 0
406
- Number of seconds to wait prior to timing out.
407
- minutes : int, default 0
408
- Number of minutes to wait prior to timing out.
409
- hours : int, default 0
410
- Number of hours to wait prior to timing out.
411
- """
412
- ...
413
-
414
- @typing.overload
415
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
454
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
416
455
  """
417
456
  Internal decorator to support Fast bakery
418
457
  """
@@ -453,55 +492,6 @@ def huggingface_hub(*, temp_dir_root: typing.Optional[str] = None, load: typing.
453
492
  """
454
493
  ...
455
494
 
456
- @typing.overload
457
- 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]]]:
458
- """
459
- Creates a human-readable report, a Metaflow Card, after this step completes.
460
-
461
- Note that you may add multiple `@card` decorators in a step with different parameters.
462
-
463
-
464
- Parameters
465
- ----------
466
- type : str, default 'default'
467
- Card type.
468
- id : str, optional, default None
469
- If multiple cards are present, use this id to identify this card.
470
- options : Dict[str, Any], default {}
471
- Options passed to the card. The contents depend on the card type.
472
- timeout : int, default 45
473
- Interrupt reporting if it takes more than this many seconds.
474
- """
475
- ...
476
-
477
- @typing.overload
478
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
479
- ...
480
-
481
- @typing.overload
482
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
483
- ...
484
-
485
- 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):
486
- """
487
- Creates a human-readable report, a Metaflow Card, after this step completes.
488
-
489
- Note that you may add multiple `@card` decorators in a step with different parameters.
490
-
491
-
492
- Parameters
493
- ----------
494
- type : str, default 'default'
495
- Card type.
496
- id : str, optional, default None
497
- If multiple cards are present, use this id to identify this card.
498
- options : Dict[str, Any], default {}
499
- Options passed to the card. The contents depend on the card type.
500
- timeout : int, default 45
501
- Interrupt reporting if it takes more than this many seconds.
502
- """
503
- ...
504
-
505
495
  @typing.overload
506
496
  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]]]:
507
497
  """
@@ -559,179 +549,146 @@ def model(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
559
549
  """
560
550
  ...
561
551
 
562
- 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]]]:
552
+ @typing.overload
553
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
563
554
  """
564
- Specifies that this step is used to deploy an instance of the app.
565
- Requires that self.app_name, self.app_port, self.entrypoint and self.deployDir is set.
566
-
567
-
568
- Parameters
569
- ----------
570
- app_port : int
571
- Number of GPUs to use.
572
- app_name : str
573
- Name of the app to deploy.
555
+ Decorator prototype for all step decorators. This function gets specialized
556
+ and imported for all decorators types by _import_plugin_decorators().
574
557
  """
575
558
  ...
576
559
 
577
- 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]]]:
560
+ @typing.overload
561
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
562
+ ...
563
+
564
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
578
565
  """
579
- This decorator is used to run Ollama APIs as Metaflow task sidecars.
580
-
581
- User code call
582
- -----------
583
- @ollama(
584
- models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
585
- backend='local'
586
- )
587
-
588
- Valid backend options
589
- ---------------------
590
- - 'local': Run as a separate process on the local task machine.
591
- - (TODO) 'managed': Outerbounds hosts and selects compute provider.
592
- - (TODO) 'remote': Spin up separate instance to serve Ollama models.
593
-
594
- Valid model options
595
- ----------------
596
- - 'llama3.2'
597
- - 'llama3.3'
598
- - any model here https://ollama.com/search
599
-
600
-
601
- Parameters
602
- ----------
603
- models: list[Ollama]
604
- List of Ollama containers running models in sidecars.
605
- backend: str
606
- Determines where and how to run the Ollama process.
566
+ Decorator prototype for all step decorators. This function gets specialized
567
+ and imported for all decorators types by _import_plugin_decorators().
607
568
  """
608
569
  ...
609
570
 
610
- def nim(*, models: "list[NIM]", backend: 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]]]:
571
+ 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]]]:
611
572
  """
612
- This decorator is used to run NIM containers in Metaflow tasks as sidecars.
613
-
614
- User code call
615
- -----------
616
- @nim(
617
- models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
618
- backend='managed'
619
- )
620
-
621
- Valid backend options
622
- ---------------------
623
- - 'managed': Outerbounds selects a compute provider based on the model.
624
-
625
- Valid model options
626
- ----------------
627
- - 'meta/llama3-8b-instruct': 8B parameter model
628
- - 'meta/llama3-70b-instruct': 70B parameter model
629
- - any model here: https://nvcf.ngc.nvidia.com/functions?filter=nvidia-functions
573
+ Specifies that this step should execute on DGX cloud.
630
574
 
631
575
 
632
576
  Parameters
633
577
  ----------
634
- models: list[NIM]
635
- List of NIM containers running models in sidecars.
636
- backend: str
637
- Compute provider to run the NIM container.
578
+ gpu : int
579
+ Number of GPUs to use.
580
+ gpu_type : str
581
+ Type of Nvidia GPU to use.
638
582
  queue_timeout : int
639
583
  Time to keep the job in NVCF's queue.
640
584
  """
641
585
  ...
642
586
 
643
587
  @typing.overload
644
- 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]]]:
588
+ 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]]]:
645
589
  """
646
- Specifies secrets to be retrieved and injected as environment variables prior to
647
- the execution of a step.
590
+ Specifies environment variables to be set prior to the execution of a step.
648
591
 
649
592
 
650
593
  Parameters
651
594
  ----------
652
- sources : List[Union[str, Dict[str, Any]]], default: []
653
- List of secret specs, defining how the secrets are to be retrieved
595
+ vars : Dict[str, str], default {}
596
+ Dictionary of environment variables to set.
654
597
  """
655
598
  ...
656
599
 
657
600
  @typing.overload
658
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
601
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
659
602
  ...
660
603
 
661
604
  @typing.overload
662
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
605
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
663
606
  ...
664
607
 
665
- 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]]] = []):
608
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
666
609
  """
667
- Specifies secrets to be retrieved and injected as environment variables prior to
668
- the execution of a step.
610
+ Specifies environment variables to be set prior to the execution of a step.
669
611
 
670
612
 
671
613
  Parameters
672
614
  ----------
673
- sources : List[Union[str, Dict[str, Any]]], default: []
674
- List of secret specs, defining how the secrets are to be retrieved
615
+ vars : Dict[str, str], default {}
616
+ Dictionary of environment variables to set.
675
617
  """
676
618
  ...
677
619
 
678
620
  @typing.overload
679
- 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]]]:
621
+ 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]]]:
680
622
  """
681
- Specifies the Conda environment for the step.
623
+ Specifies a timeout for your step.
682
624
 
683
- Information in this decorator will augment any
684
- attributes set in the `@conda_base` flow-level decorator. Hence,
685
- you can use `@conda_base` to set packages required by all
686
- steps and use `@conda` to specify step-specific overrides.
625
+ This decorator is useful if this step may hang indefinitely.
626
+
627
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
628
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
629
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
630
+
631
+ Note that all the values specified in parameters are added together so if you specify
632
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
687
633
 
688
634
 
689
635
  Parameters
690
636
  ----------
691
- packages : Dict[str, str], default {}
692
- Packages to use for this step. The key is the name of the package
693
- and the value is the version to use.
694
- libraries : Dict[str, str], default {}
695
- Supported for backward compatibility. When used with packages, packages will take precedence.
696
- python : str, optional, default None
697
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
698
- that the version used will correspond to the version of the Python interpreter used to start the run.
699
- disabled : bool, default False
700
- If set to True, disables @conda.
637
+ seconds : int, default 0
638
+ Number of seconds to wait prior to timing out.
639
+ minutes : int, default 0
640
+ Number of minutes to wait prior to timing out.
641
+ hours : int, default 0
642
+ Number of hours to wait prior to timing out.
701
643
  """
702
644
  ...
703
645
 
704
646
  @typing.overload
705
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
647
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
706
648
  ...
707
649
 
708
650
  @typing.overload
709
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
651
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
710
652
  ...
711
653
 
712
- 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):
654
+ 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):
713
655
  """
714
- Specifies the Conda environment for the step.
656
+ Specifies a timeout for your step.
715
657
 
716
- Information in this decorator will augment any
717
- attributes set in the `@conda_base` flow-level decorator. Hence,
718
- you can use `@conda_base` to set packages required by all
719
- steps and use `@conda` to specify step-specific overrides.
658
+ This decorator is useful if this step may hang indefinitely.
659
+
660
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
661
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
662
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
663
+
664
+ Note that all the values specified in parameters are added together so if you specify
665
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
720
666
 
721
667
 
722
668
  Parameters
723
669
  ----------
724
- packages : Dict[str, str], default {}
725
- Packages to use for this step. The key is the name of the package
726
- and the value is the version to use.
727
- libraries : Dict[str, str], default {}
728
- Supported for backward compatibility. When used with packages, packages will take precedence.
729
- python : str, optional, default None
730
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
731
- that the version used will correspond to the version of the Python interpreter used to start the run.
732
- disabled : bool, default False
733
- If set to True, disables @conda.
734
- """
670
+ seconds : int, default 0
671
+ Number of seconds to wait prior to timing out.
672
+ minutes : int, default 0
673
+ Number of minutes to wait prior to timing out.
674
+ hours : int, default 0
675
+ Number of hours to wait prior to timing out.
676
+ """
677
+ ...
678
+
679
+ 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]]]:
680
+ """
681
+ Specifies that this step is used to deploy an instance of the app.
682
+ Requires that self.app_name, self.app_port, self.entrypoint and self.deployDir is set.
683
+
684
+
685
+ Parameters
686
+ ----------
687
+ app_port : int
688
+ Number of GPUs to use.
689
+ app_name : str
690
+ Name of the app to deploy.
691
+ """
735
692
  ...
736
693
 
737
694
  @typing.overload
@@ -786,54 +743,175 @@ def pypi(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typ
786
743
  ...
787
744
 
788
745
  @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]]]:
746
+ 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]]]:
790
747
  """
791
- Specifies environment variables to be set prior to the execution of a step.
748
+ Specifies the resources needed when executing this step.
749
+
750
+ Use `@resources` to specify the resource requirements
751
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
752
+
753
+ You can choose the compute layer on the command line by executing e.g.
754
+ ```
755
+ python myflow.py run --with batch
756
+ ```
757
+ or
758
+ ```
759
+ python myflow.py run --with kubernetes
760
+ ```
761
+ which executes the flow on the desired system using the
762
+ requirements specified in `@resources`.
792
763
 
793
764
 
794
765
  Parameters
795
766
  ----------
796
- vars : Dict[str, str], default {}
797
- Dictionary of environment variables to set.
767
+ cpu : int, default 1
768
+ Number of CPUs required for this step.
769
+ gpu : int, optional, default None
770
+ Number of GPUs required for this step.
771
+ disk : int, optional, default None
772
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
773
+ memory : int, default 4096
774
+ Memory size (in MB) required for this step.
775
+ shared_memory : int, optional, default None
776
+ The value for the size (in MiB) of the /dev/shm volume for this step.
777
+ This parameter maps to the `--shm-size` option in Docker.
798
778
  """
799
779
  ...
800
780
 
801
781
  @typing.overload
802
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
782
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
803
783
  ...
804
784
 
805
785
  @typing.overload
806
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
786
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
807
787
  ...
808
788
 
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] = {}):
789
+ 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):
810
790
  """
811
- Specifies environment variables to be set prior to the execution of a step.
791
+ Specifies the resources needed when executing this step.
792
+
793
+ Use `@resources` to specify the resource requirements
794
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
795
+
796
+ You can choose the compute layer on the command line by executing e.g.
797
+ ```
798
+ python myflow.py run --with batch
799
+ ```
800
+ or
801
+ ```
802
+ python myflow.py run --with kubernetes
803
+ ```
804
+ which executes the flow on the desired system using the
805
+ requirements specified in `@resources`.
812
806
 
813
807
 
814
808
  Parameters
815
809
  ----------
816
- vars : Dict[str, str], default {}
817
- Dictionary of environment variables to set.
810
+ cpu : int, default 1
811
+ Number of CPUs required for this step.
812
+ gpu : int, optional, default None
813
+ Number of GPUs required for this step.
814
+ disk : int, optional, default None
815
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
816
+ memory : int, default 4096
817
+ Memory size (in MB) required for this step.
818
+ shared_memory : int, optional, default None
819
+ The value for the size (in MiB) of the /dev/shm volume for this step.
820
+ This parameter maps to the `--shm-size` option in Docker.
818
821
  """
819
822
  ...
820
823
 
821
824
  @typing.overload
822
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
825
+ 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]]]:
823
826
  """
824
- Decorator prototype for all step decorators. This function gets specialized
825
- and imported for all decorators types by _import_plugin_decorators().
827
+ Specifies the Conda environment for the step.
828
+
829
+ Information in this decorator will augment any
830
+ attributes set in the `@conda_base` flow-level decorator. Hence,
831
+ you can use `@conda_base` to set packages required by all
832
+ steps and use `@conda` to specify step-specific overrides.
833
+
834
+
835
+ Parameters
836
+ ----------
837
+ packages : Dict[str, str], default {}
838
+ Packages to use for this step. The key is the name of the package
839
+ and the value is the version to use.
840
+ libraries : Dict[str, str], default {}
841
+ Supported for backward compatibility. When used with packages, packages will take precedence.
842
+ python : str, optional, default None
843
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
844
+ that the version used will correspond to the version of the Python interpreter used to start the run.
845
+ disabled : bool, default False
846
+ If set to True, disables @conda.
826
847
  """
827
848
  ...
828
849
 
829
850
  @typing.overload
830
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
851
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
831
852
  ...
832
853
 
833
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
854
+ @typing.overload
855
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
856
+ ...
857
+
858
+ 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):
834
859
  """
835
- Decorator prototype for all step decorators. This function gets specialized
836
- and imported for all decorators types by _import_plugin_decorators().
860
+ Specifies the Conda environment for the step.
861
+
862
+ Information in this decorator will augment any
863
+ attributes set in the `@conda_base` flow-level decorator. Hence,
864
+ you can use `@conda_base` to set packages required by all
865
+ steps and use `@conda` to specify step-specific overrides.
866
+
867
+
868
+ Parameters
869
+ ----------
870
+ packages : Dict[str, str], default {}
871
+ Packages to use for this step. The key is the name of the package
872
+ and the value is the version to use.
873
+ libraries : Dict[str, str], default {}
874
+ Supported for backward compatibility. When used with packages, packages will take precedence.
875
+ python : str, optional, default None
876
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
877
+ that the version used will correspond to the version of the Python interpreter used to start the run.
878
+ disabled : bool, default False
879
+ If set to True, disables @conda.
880
+ """
881
+ ...
882
+
883
+ @typing.overload
884
+ 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]]]:
885
+ """
886
+ Specifies secrets to be retrieved and injected as environment variables prior to
887
+ the execution of a step.
888
+
889
+
890
+ Parameters
891
+ ----------
892
+ sources : List[Union[str, Dict[str, Any]]], default: []
893
+ List of secret specs, defining how the secrets are to be retrieved
894
+ """
895
+ ...
896
+
897
+ @typing.overload
898
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
899
+ ...
900
+
901
+ @typing.overload
902
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
903
+ ...
904
+
905
+ 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]]] = []):
906
+ """
907
+ Specifies secrets to be retrieved and injected as environment variables prior to
908
+ the execution of a step.
909
+
910
+
911
+ Parameters
912
+ ----------
913
+ sources : List[Union[str, Dict[str, Any]]], default: []
914
+ List of secret specs, defining how the secrets are to be retrieved
837
915
  """
838
916
  ...
839
917
 
@@ -892,80 +970,104 @@ def retry(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
892
970
  """
893
971
  ...
894
972
 
895
- def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: typing.Optional[str] = None, image_pull_policy: str = 'KUBERNETES_IMAGE_PULL_POLICY', service_account: str = 'METAFLOW_KUBERNETES_SERVICE_ACCOUNT', secrets: typing.Optional[typing.List[str]] = None, node_selector: typing.Union[typing.Dict[str, str], str, None] = None, namespace: str = 'METAFLOW_KUBERNETES_NAMESPACE', gpu: typing.Optional[int] = None, gpu_vendor: str = 'KUBERNETES_GPU_VENDOR', tolerations: typing.List[str] = [], labels: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_LABELS', annotations: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_ANNOTATIONS', use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = '/metaflow_temp', persistent_volume_claims: typing.Optional[typing.Dict[str, str]] = None, shared_memory: typing.Optional[int] = None, port: typing.Optional[int] = None, compute_pool: typing.Optional[str] = None, hostname_resolution_timeout: int = 600, qos: str = 'Burstable') -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
973
+ @typing.overload
974
+ 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]]:
896
975
  """
897
- Specifies that this step should execute on Kubernetes.
976
+ Specifies the flow(s) that this flow depends on.
977
+
978
+ ```
979
+ @trigger_on_finish(flow='FooFlow')
980
+ ```
981
+ or
982
+ ```
983
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
984
+ ```
985
+ This decorator respects the @project decorator and triggers the flow
986
+ when upstream runs within the same namespace complete successfully
987
+
988
+ Additionally, you can specify project aware upstream flow dependencies
989
+ by specifying the fully qualified project_flow_name.
990
+ ```
991
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
992
+ ```
993
+ or
994
+ ```
995
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
996
+ ```
997
+
998
+ You can also specify just the project or project branch (other values will be
999
+ inferred from the current project or project branch):
1000
+ ```
1001
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1002
+ ```
1003
+
1004
+ Note that `branch` is typically one of:
1005
+ - `prod`
1006
+ - `user.bob`
1007
+ - `test.my_experiment`
1008
+ - `prod.staging`
898
1009
 
899
1010
 
900
1011
  Parameters
901
1012
  ----------
902
- cpu : int, default 1
903
- Number of CPUs required for this step. If `@resources` is
904
- also present, the maximum value from all decorators is used.
905
- memory : int, default 4096
906
- Memory size (in MB) required for this step. If
907
- `@resources` is also present, the maximum value from all decorators is
908
- used.
909
- disk : int, default 10240
910
- Disk size (in MB) required for this step. If
911
- `@resources` is also present, the maximum value from all decorators is
912
- used.
913
- image : str, optional, default None
914
- Docker image to use when launching on Kubernetes. If not specified, and
915
- METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
916
- not, a default Docker image mapping to the current version of Python is used.
917
- image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
918
- If given, the imagePullPolicy to be applied to the Docker image of the step.
919
- service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
920
- Kubernetes service account to use when launching pod in Kubernetes.
921
- secrets : List[str], optional, default None
922
- Kubernetes secrets to use when launching pod in Kubernetes. These
923
- secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
924
- in Metaflow configuration.
925
- node_selector: Union[Dict[str,str], str], optional, default None
926
- Kubernetes node selector(s) to apply to the pod running the task.
927
- Can be passed in as a comma separated string of values e.g.
928
- 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
929
- {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
930
- namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
931
- Kubernetes namespace to use when launching pod in Kubernetes.
932
- gpu : int, optional, default None
933
- Number of GPUs required for this step. A value of zero implies that
934
- the scheduled node should not have GPUs.
935
- gpu_vendor : str, default KUBERNETES_GPU_VENDOR
936
- The vendor of the GPUs to be used for this step.
937
- tolerations : List[str], default []
938
- The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
939
- Kubernetes tolerations to use when launching pod in Kubernetes.
940
- labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
941
- Kubernetes labels to use when launching pod in Kubernetes.
942
- annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
943
- Kubernetes annotations to use when launching pod in Kubernetes.
944
- use_tmpfs : bool, default False
945
- This enables an explicit tmpfs mount for this step.
946
- tmpfs_tempdir : bool, default True
947
- sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
948
- tmpfs_size : int, optional, default: None
949
- The value for the size (in MiB) of the tmpfs mount for this step.
950
- This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
951
- memory allocated for this step.
952
- tmpfs_path : str, optional, default /metaflow_temp
953
- Path to tmpfs mount for this step.
954
- persistent_volume_claims : Dict[str, str], optional, default None
955
- A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
956
- volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
957
- shared_memory: int, optional
958
- Shared memory size (in MiB) required for this step
959
- port: int, optional
960
- Port number to specify in the Kubernetes job object
961
- compute_pool : str, optional, default None
962
- Compute pool to be used for for this step.
963
- If not specified, any accessible compute pool within the perimeter is used.
964
- hostname_resolution_timeout: int, default 10 * 60
965
- Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
966
- Only applicable when @parallel is used.
967
- qos: str, default: Burstable
968
- Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
1013
+ flow : Union[str, Dict[str, str]], optional, default None
1014
+ Upstream flow dependency for this flow.
1015
+ flows : List[Union[str, Dict[str, str]]], default []
1016
+ Upstream flow dependencies for this flow.
1017
+ options : Dict[str, Any], default {}
1018
+ Backend-specific configuration for tuning eventing behavior.
1019
+ """
1020
+ ...
1021
+
1022
+ @typing.overload
1023
+ def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1024
+ ...
1025
+
1026
+ def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}):
1027
+ """
1028
+ Specifies the flow(s) that this flow depends on.
1029
+
1030
+ ```
1031
+ @trigger_on_finish(flow='FooFlow')
1032
+ ```
1033
+ or
1034
+ ```
1035
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1036
+ ```
1037
+ This decorator respects the @project decorator and triggers the flow
1038
+ when upstream runs within the same namespace complete successfully
1039
+
1040
+ Additionally, you can specify project aware upstream flow dependencies
1041
+ by specifying the fully qualified project_flow_name.
1042
+ ```
1043
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1044
+ ```
1045
+ or
1046
+ ```
1047
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1048
+ ```
1049
+
1050
+ You can also specify just the project or project branch (other values will be
1051
+ inferred from the current project or project branch):
1052
+ ```
1053
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1054
+ ```
1055
+
1056
+ Note that `branch` is typically one of:
1057
+ - `prod`
1058
+ - `user.bob`
1059
+ - `test.my_experiment`
1060
+ - `prod.staging`
1061
+
1062
+
1063
+ Parameters
1064
+ ----------
1065
+ flow : Union[str, Dict[str, str]], optional, default None
1066
+ Upstream flow dependency for this flow.
1067
+ flows : List[Union[str, Dict[str, str]]], default []
1068
+ Upstream flow dependencies for this flow.
1069
+ options : Dict[str, Any], default {}
1070
+ Backend-specific configuration for tuning eventing behavior.
969
1071
  """
970
1072
  ...
971
1073
 
@@ -1010,54 +1112,203 @@ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packag
1010
1112
  """
1011
1113
  ...
1012
1114
 
1013
- @typing.overload
1014
- 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]]:
1115
+ 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]]:
1015
1116
  """
1016
- Specifies the Conda environment for all steps of the flow.
1017
-
1018
- Use `@conda_base` to set common libraries required by all
1019
- steps and use `@conda` to specify step-specific additions.
1117
+ 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)
1118
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1119
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1120
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
1121
+ starts only after all sensors finish.
1020
1122
 
1021
1123
 
1022
1124
  Parameters
1023
1125
  ----------
1024
- packages : Dict[str, str], default {}
1025
- Packages to use for this flow. The key is the name of the package
1026
- and the value is the version to use.
1027
- libraries : Dict[str, str], default {}
1028
- Supported for backward compatibility. When used with packages, packages will take precedence.
1029
- python : str, optional, default None
1030
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1031
- that the version used will correspond to the version of the Python interpreter used to start the run.
1032
- disabled : bool, default False
1033
- If set to True, disables Conda.
1126
+ timeout : int
1127
+ Time, in seconds before the task times out and fails. (Default: 3600)
1128
+ poke_interval : int
1129
+ Time in seconds that the job should wait in between each try. (Default: 60)
1130
+ mode : str
1131
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1132
+ exponential_backoff : bool
1133
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1134
+ pool : str
1135
+ the slot pool this task should run in,
1136
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1137
+ soft_fail : bool
1138
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1139
+ name : str
1140
+ Name of the sensor on Airflow
1141
+ description : str
1142
+ Description of sensor in the Airflow UI
1143
+ bucket_key : Union[str, List[str]]
1144
+ The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1145
+ When it's specified as a full s3:// url, please leave `bucket_name` as None
1146
+ bucket_name : str
1147
+ Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1148
+ When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1149
+ wildcard_match : bool
1150
+ whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1151
+ aws_conn_id : str
1152
+ a reference to the s3 connection on Airflow. (Default: None)
1153
+ verify : bool
1154
+ Whether or not to verify SSL certificates for S3 connection. (Default: None)
1034
1155
  """
1035
1156
  ...
1036
1157
 
1037
- @typing.overload
1038
- def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1158
+ def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None):
1159
+ """
1160
+ Allows setting external datastores to save data for the
1161
+ `@checkpoint`/`@model`/`@huggingface_hub` decorators.
1162
+
1163
+ This decorator is useful when users wish to save data to a different datastore
1164
+ than what is configured in Metaflow. This can be for variety of reasons:
1165
+
1166
+ 1. Data security: The objects needs to be stored in a bucket (object storage) that is not accessible by other flows.
1167
+ 2. Data Locality: The location where the task is executing is not located in the same region as the datastore.
1168
+ - Example: Metaflow datastore lives in US East, but the task is executing in Finland datacenters.
1169
+ 3. Data Lifecycle Policies: The objects need to be archived / managed separately from the Metaflow managed objects.
1170
+ - Example: Flow is training very large models that need to be stored separately and will be deleted more aggressively than the Metaflow managed objects.
1171
+
1172
+ Usage:
1173
+ ----------
1174
+
1175
+ - Using a custom IAM role to access the datastore.
1176
+
1177
+ ```python
1178
+ @with_artifact_store(
1179
+ type="s3",
1180
+ config=lambda: {
1181
+ "root": "s3://my-bucket-foo/path/to/root",
1182
+ "role_arn": ROLE,
1183
+ },
1184
+ )
1185
+ class MyFlow(FlowSpec):
1186
+
1187
+ @checkpoint
1188
+ @step
1189
+ def start(self):
1190
+ with open("my_file.txt", "w") as f:
1191
+ f.write("Hello, World!")
1192
+ self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1193
+ self.next(self.end)
1194
+
1195
+ ```
1196
+
1197
+ - Using credentials to access the s3-compatible datastore.
1198
+
1199
+ ```python
1200
+ @with_artifact_store(
1201
+ type="s3",
1202
+ config=lambda: {
1203
+ "root": "s3://my-bucket-foo/path/to/root",
1204
+ "client_params": {
1205
+ "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1206
+ "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1207
+ },
1208
+ },
1209
+ )
1210
+ class MyFlow(FlowSpec):
1211
+
1212
+ @checkpoint
1213
+ @step
1214
+ def start(self):
1215
+ with open("my_file.txt", "w") as f:
1216
+ f.write("Hello, World!")
1217
+ self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1218
+ self.next(self.end)
1219
+
1220
+ ```
1221
+
1222
+ - Accessing objects stored in external datastores after task execution.
1223
+
1224
+ ```python
1225
+ run = Run("CheckpointsTestsFlow/8992")
1226
+ with artifact_store_from(run=run, config={
1227
+ "client_params": {
1228
+ "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1229
+ "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1230
+ },
1231
+ }):
1232
+ with Checkpoint() as cp:
1233
+ latest = cp.list(
1234
+ task=run["start"].task
1235
+ )[0]
1236
+ print(latest)
1237
+ cp.load(
1238
+ latest,
1239
+ "test-checkpoints"
1240
+ )
1241
+
1242
+ task = Task("TorchTuneFlow/8484/train/53673")
1243
+ with artifact_store_from(run=run, config={
1244
+ "client_params": {
1245
+ "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1246
+ "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1247
+ },
1248
+ }):
1249
+ load_model(
1250
+ task.data.model_ref,
1251
+ "test-models"
1252
+ )
1253
+ ```
1254
+ Parameters:
1255
+ ----------
1256
+
1257
+ type: str
1258
+ The type of the datastore. Can be one of 's3', 'gcs', 'azure' or any other supported metaflow Datastore.
1259
+
1260
+ config: dict or Callable
1261
+ Dictionary of configuration options for the datastore. The following keys are required:
1262
+ - root: The root path in the datastore where the data will be saved. (needs to be in the format expected by the datastore)
1263
+ - example: 's3://bucket-name/path/to/root'
1264
+ - example: 'gs://bucket-name/path/to/root'
1265
+ - example: 'https://myblockacc.blob.core.windows.net/metaflow/'
1266
+ - role_arn (optional): AWS IAM role to access s3 bucket (only when `type` is 's3')
1267
+ - session_vars (optional): AWS session variables to access s3 bucket (only when `type` is 's3')
1268
+ - client_params (optional): AWS client parameters to access s3 bucket (only when `type` is 's3')
1269
+ """
1039
1270
  ...
1040
1271
 
1041
- 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):
1272
+ 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]]:
1042
1273
  """
1043
- Specifies the Conda environment for all steps of the flow.
1044
-
1045
- Use `@conda_base` to set common libraries required by all
1046
- steps and use `@conda` to specify step-specific additions.
1274
+ 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.
1275
+ 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.
1047
1276
 
1048
1277
 
1049
1278
  Parameters
1050
1279
  ----------
1051
- packages : Dict[str, str], default {}
1052
- Packages to use for this flow. The key is the name of the package
1053
- and the value is the version to use.
1054
- libraries : Dict[str, str], default {}
1055
- Supported for backward compatibility. When used with packages, packages will take precedence.
1056
- python : str, optional, default None
1057
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1058
- that the version used will correspond to the version of the Python interpreter used to start the run.
1059
- disabled : bool, default False
1060
- If set to True, disables Conda.
1280
+ timeout : int
1281
+ Time, in seconds before the task times out and fails. (Default: 3600)
1282
+ poke_interval : int
1283
+ Time in seconds that the job should wait in between each try. (Default: 60)
1284
+ mode : str
1285
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1286
+ exponential_backoff : bool
1287
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1288
+ pool : str
1289
+ the slot pool this task should run in,
1290
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1291
+ soft_fail : bool
1292
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1293
+ name : str
1294
+ Name of the sensor on Airflow
1295
+ description : str
1296
+ Description of sensor in the Airflow UI
1297
+ external_dag_id : str
1298
+ The dag_id that contains the task you want to wait for.
1299
+ external_task_ids : List[str]
1300
+ The list of task_ids that you want to wait for.
1301
+ If None (default value) the sensor waits for the DAG. (Default: None)
1302
+ allowed_states : List[str]
1303
+ Iterable of allowed states, (Default: ['success'])
1304
+ failed_states : List[str]
1305
+ Iterable of failed or dis-allowed states. (Default: None)
1306
+ execution_delta : datetime.timedelta
1307
+ time difference with the previous execution to look at,
1308
+ the default is the same logical date as the current task or DAG. (Default: None)
1309
+ check_existence: bool
1310
+ Set to True to check if the external task exists or check if
1311
+ the DAG to wait for exists. (Default: True)
1061
1312
  """
1062
1313
  ...
1063
1314
 
@@ -1096,6 +1347,57 @@ def project(*, name: str, branch: typing.Optional[str] = None, production: bool
1096
1347
  """
1097
1348
  ...
1098
1349
 
1350
+ @typing.overload
1351
+ 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]]:
1352
+ """
1353
+ Specifies the times when the flow should be run when running on a
1354
+ production scheduler.
1355
+
1356
+
1357
+ Parameters
1358
+ ----------
1359
+ hourly : bool, default False
1360
+ Run the workflow hourly.
1361
+ daily : bool, default True
1362
+ Run the workflow daily.
1363
+ weekly : bool, default False
1364
+ Run the workflow weekly.
1365
+ cron : str, optional, default None
1366
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1367
+ specified by this expression.
1368
+ timezone : str, optional, default None
1369
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1370
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1371
+ """
1372
+ ...
1373
+
1374
+ @typing.overload
1375
+ def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1376
+ ...
1377
+
1378
+ def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None):
1379
+ """
1380
+ Specifies the times when the flow should be run when running on a
1381
+ production scheduler.
1382
+
1383
+
1384
+ Parameters
1385
+ ----------
1386
+ hourly : bool, default False
1387
+ Run the workflow hourly.
1388
+ daily : bool, default True
1389
+ Run the workflow daily.
1390
+ weekly : bool, default False
1391
+ Run the workflow weekly.
1392
+ cron : str, optional, default None
1393
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1394
+ specified by this expression.
1395
+ timezone : str, optional, default None
1396
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1397
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1398
+ """
1399
+ ...
1400
+
1099
1401
  @typing.overload
1100
1402
  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]]:
1101
1403
  """
@@ -1189,241 +1491,54 @@ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: t
1189
1491
  """
1190
1492
  ...
1191
1493
 
1192
- 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]]:
1193
- """
1194
- 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.
1195
- 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.
1196
-
1197
-
1198
- Parameters
1199
- ----------
1200
- timeout : int
1201
- Time, in seconds before the task times out and fails. (Default: 3600)
1202
- poke_interval : int
1203
- Time in seconds that the job should wait in between each try. (Default: 60)
1204
- mode : str
1205
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1206
- exponential_backoff : bool
1207
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1208
- pool : str
1209
- the slot pool this task should run in,
1210
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1211
- soft_fail : bool
1212
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1213
- name : str
1214
- Name of the sensor on Airflow
1215
- description : str
1216
- Description of sensor in the Airflow UI
1217
- external_dag_id : str
1218
- The dag_id that contains the task you want to wait for.
1219
- external_task_ids : List[str]
1220
- The list of task_ids that you want to wait for.
1221
- If None (default value) the sensor waits for the DAG. (Default: None)
1222
- allowed_states : List[str]
1223
- Iterable of allowed states, (Default: ['success'])
1224
- failed_states : List[str]
1225
- Iterable of failed or dis-allowed states. (Default: None)
1226
- execution_delta : datetime.timedelta
1227
- time difference with the previous execution to look at,
1228
- the default is the same logical date as the current task or DAG. (Default: None)
1229
- check_existence: bool
1230
- Set to True to check if the external task exists or check if
1231
- the DAG to wait for exists. (Default: True)
1232
- """
1233
- ...
1234
-
1235
1494
  @typing.overload
1236
- 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]]:
1495
+ 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]]:
1237
1496
  """
1238
- Specifies the flow(s) that this flow depends on.
1239
-
1240
- ```
1241
- @trigger_on_finish(flow='FooFlow')
1242
- ```
1243
- or
1244
- ```
1245
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1246
- ```
1247
- This decorator respects the @project decorator and triggers the flow
1248
- when upstream runs within the same namespace complete successfully
1249
-
1250
- Additionally, you can specify project aware upstream flow dependencies
1251
- by specifying the fully qualified project_flow_name.
1252
- ```
1253
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1254
- ```
1255
- or
1256
- ```
1257
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1258
- ```
1259
-
1260
- You can also specify just the project or project branch (other values will be
1261
- inferred from the current project or project branch):
1262
- ```
1263
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1264
- ```
1497
+ Specifies the Conda environment for all steps of the flow.
1265
1498
 
1266
- Note that `branch` is typically one of:
1267
- - `prod`
1268
- - `user.bob`
1269
- - `test.my_experiment`
1270
- - `prod.staging`
1499
+ Use `@conda_base` to set common libraries required by all
1500
+ steps and use `@conda` to specify step-specific additions.
1271
1501
 
1272
1502
 
1273
1503
  Parameters
1274
1504
  ----------
1275
- flow : Union[str, Dict[str, str]], optional, default None
1276
- Upstream flow dependency for this flow.
1277
- flows : List[Union[str, Dict[str, str]]], default []
1278
- Upstream flow dependencies for this flow.
1279
- options : Dict[str, Any], default {}
1280
- Backend-specific configuration for tuning eventing behavior.
1505
+ packages : Dict[str, str], default {}
1506
+ Packages to use for this flow. The key is the name of the package
1507
+ and the value is the version to use.
1508
+ libraries : Dict[str, str], default {}
1509
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1510
+ python : str, optional, default None
1511
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1512
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1513
+ disabled : bool, default False
1514
+ If set to True, disables Conda.
1281
1515
  """
1282
1516
  ...
1283
1517
 
1284
1518
  @typing.overload
1285
- def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1286
- ...
1287
-
1288
- def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}):
1289
- """
1290
- Specifies the flow(s) that this flow depends on.
1291
-
1292
- ```
1293
- @trigger_on_finish(flow='FooFlow')
1294
- ```
1295
- or
1296
- ```
1297
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1298
- ```
1299
- This decorator respects the @project decorator and triggers the flow
1300
- when upstream runs within the same namespace complete successfully
1301
-
1302
- Additionally, you can specify project aware upstream flow dependencies
1303
- by specifying the fully qualified project_flow_name.
1304
- ```
1305
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1306
- ```
1307
- or
1308
- ```
1309
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1310
- ```
1311
-
1312
- You can also specify just the project or project branch (other values will be
1313
- inferred from the current project or project branch):
1314
- ```
1315
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1316
- ```
1317
-
1318
- Note that `branch` is typically one of:
1319
- - `prod`
1320
- - `user.bob`
1321
- - `test.my_experiment`
1322
- - `prod.staging`
1323
-
1324
-
1325
- Parameters
1326
- ----------
1327
- flow : Union[str, Dict[str, str]], optional, default None
1328
- Upstream flow dependency for this flow.
1329
- flows : List[Union[str, Dict[str, str]]], default []
1330
- Upstream flow dependencies for this flow.
1331
- options : Dict[str, Any], default {}
1332
- Backend-specific configuration for tuning eventing behavior.
1333
- """
1334
- ...
1335
-
1336
- 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]]:
1337
- """
1338
- 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)
1339
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1340
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1341
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
1342
- starts only after all sensors finish.
1343
-
1344
-
1345
- Parameters
1346
- ----------
1347
- timeout : int
1348
- Time, in seconds before the task times out and fails. (Default: 3600)
1349
- poke_interval : int
1350
- Time in seconds that the job should wait in between each try. (Default: 60)
1351
- mode : str
1352
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1353
- exponential_backoff : bool
1354
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1355
- pool : str
1356
- the slot pool this task should run in,
1357
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1358
- soft_fail : bool
1359
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1360
- name : str
1361
- Name of the sensor on Airflow
1362
- description : str
1363
- Description of sensor in the Airflow UI
1364
- bucket_key : Union[str, List[str]]
1365
- The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1366
- When it's specified as a full s3:// url, please leave `bucket_name` as None
1367
- bucket_name : str
1368
- Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1369
- When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1370
- wildcard_match : bool
1371
- whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1372
- aws_conn_id : str
1373
- a reference to the s3 connection on Airflow. (Default: None)
1374
- verify : bool
1375
- Whether or not to verify SSL certificates for S3 connection. (Default: None)
1376
- """
1519
+ def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1377
1520
  ...
1378
1521
 
1379
- @typing.overload
1380
- 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]]:
1522
+ 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):
1381
1523
  """
1382
- Specifies the times when the flow should be run when running on a
1383
- production scheduler.
1384
-
1524
+ Specifies the Conda environment for all steps of the flow.
1385
1525
 
1386
- Parameters
1387
- ----------
1388
- hourly : bool, default False
1389
- Run the workflow hourly.
1390
- daily : bool, default True
1391
- Run the workflow daily.
1392
- weekly : bool, default False
1393
- Run the workflow weekly.
1394
- cron : str, optional, default None
1395
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1396
- specified by this expression.
1397
- timezone : str, optional, default None
1398
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1399
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1400
- """
1401
- ...
1402
-
1403
- @typing.overload
1404
- def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1405
- ...
1406
-
1407
- def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None):
1408
- """
1409
- Specifies the times when the flow should be run when running on a
1410
- production scheduler.
1526
+ Use `@conda_base` to set common libraries required by all
1527
+ steps and use `@conda` to specify step-specific additions.
1411
1528
 
1412
1529
 
1413
1530
  Parameters
1414
1531
  ----------
1415
- hourly : bool, default False
1416
- Run the workflow hourly.
1417
- daily : bool, default True
1418
- Run the workflow daily.
1419
- weekly : bool, default False
1420
- Run the workflow weekly.
1421
- cron : str, optional, default None
1422
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1423
- specified by this expression.
1424
- timezone : str, optional, default None
1425
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1426
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1532
+ packages : Dict[str, str], default {}
1533
+ Packages to use for this flow. The key is the name of the package
1534
+ and the value is the version to use.
1535
+ libraries : Dict[str, str], default {}
1536
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1537
+ python : str, optional, default None
1538
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1539
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1540
+ disabled : bool, default False
1541
+ If set to True, disables Conda.
1427
1542
  """
1428
1543
  ...
1429
1544