ob-metaflow-stubs 6.0.3.184__py2.py3-none-any.whl → 6.0.3.185__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 (216) hide show
  1. metaflow-stubs/__init__.pyi +468 -468
  2. metaflow-stubs/cards.pyi +1 -1
  3. metaflow-stubs/cli.pyi +1 -1
  4. metaflow-stubs/cli_components/__init__.pyi +1 -1
  5. metaflow-stubs/cli_components/utils.pyi +1 -1
  6. metaflow-stubs/client/__init__.pyi +1 -1
  7. metaflow-stubs/client/core.pyi +4 -4
  8. metaflow-stubs/client/filecache.pyi +2 -2
  9. metaflow-stubs/events.pyi +1 -1
  10. metaflow-stubs/exception.pyi +1 -1
  11. metaflow-stubs/flowspec.pyi +3 -3
  12. metaflow-stubs/generated_for.txt +1 -1
  13. metaflow-stubs/includefile.pyi +2 -2
  14. metaflow-stubs/info_file.pyi +1 -1
  15. metaflow-stubs/metadata_provider/__init__.pyi +1 -1
  16. metaflow-stubs/metadata_provider/heartbeat.pyi +1 -1
  17. metaflow-stubs/metadata_provider/metadata.pyi +1 -1
  18. metaflow-stubs/metadata_provider/util.pyi +1 -1
  19. metaflow-stubs/metaflow_config.pyi +1 -1
  20. metaflow-stubs/metaflow_current.pyi +98 -98
  21. metaflow-stubs/metaflow_git.pyi +1 -1
  22. metaflow-stubs/mf_extensions/__init__.pyi +1 -1
  23. metaflow-stubs/mf_extensions/obcheckpoint/__init__.pyi +1 -1
  24. metaflow-stubs/mf_extensions/obcheckpoint/plugins/__init__.pyi +1 -1
  25. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/__init__.pyi +1 -1
  26. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/__init__.pyi +1 -1
  27. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/async_cards.pyi +2 -2
  28. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/deco_injection_mixin.pyi +1 -1
  29. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/extra_components.pyi +1 -1
  30. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/__init__.pyi +1 -1
  31. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/__init__.pyi +1 -1
  32. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/checkpoint_lister.pyi +3 -3
  33. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/lineage_card.pyi +1 -1
  34. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/checkpoint_storage.pyi +3 -3
  35. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/constructors.pyi +1 -1
  36. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/core.pyi +3 -3
  37. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/decorator.pyi +3 -3
  38. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/exceptions.pyi +1 -1
  39. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/final_api.pyi +1 -1
  40. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/lineage.pyi +1 -1
  41. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/__init__.pyi +1 -1
  42. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/context.pyi +1 -1
  43. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/core.pyi +1 -1
  44. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/decorator.pyi +1 -1
  45. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/exceptions.pyi +1 -1
  46. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/task_utils.pyi +1 -1
  47. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/utils.pyi +1 -1
  48. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastructures.pyi +1 -1
  49. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/exceptions.pyi +1 -1
  50. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/__init__.pyi +1 -1
  51. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/decorator.pyi +2 -2
  52. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/__init__.pyi +1 -1
  53. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/core.pyi +1 -1
  54. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/exceptions.pyi +1 -1
  55. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/model_storage.pyi +2 -2
  56. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/__init__.pyi +1 -1
  57. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/flowspec_utils.pyi +1 -1
  58. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/general.pyi +1 -1
  59. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/identity_utils.pyi +1 -1
  60. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/__init__.pyi +1 -1
  61. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/base.pyi +1 -1
  62. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/tar.pyi +1 -1
  63. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/tar_utils.pyi +1 -1
  64. metaflow-stubs/mf_extensions/outerbounds/__init__.pyi +1 -1
  65. metaflow-stubs/mf_extensions/outerbounds/plugins/__init__.pyi +1 -1
  66. metaflow-stubs/mf_extensions/outerbounds/plugins/aws/__init__.pyi +1 -1
  67. metaflow-stubs/mf_extensions/outerbounds/plugins/aws/assume_role_decorator.pyi +1 -1
  68. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/__init__.pyi +1 -1
  69. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/async_cards.pyi +1 -1
  70. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/injector.pyi +1 -1
  71. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/__init__.pyi +1 -1
  72. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/coreweave.pyi +1 -1
  73. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/nebius.pyi +1 -1
  74. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/__init__.pyi +1 -1
  75. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/pod_killer.pyi +1 -1
  76. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/__init__.pyi +1 -1
  77. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/constants.pyi +1 -1
  78. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/exceptions.pyi +1 -1
  79. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/ollama.pyi +1 -1
  80. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/status_card.pyi +1 -1
  81. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/__init__.pyi +1 -1
  82. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/snowflake.pyi +1 -1
  83. metaflow-stubs/mf_extensions/outerbounds/profilers/__init__.pyi +1 -1
  84. metaflow-stubs/mf_extensions/outerbounds/profilers/gpu.pyi +1 -1
  85. metaflow-stubs/mf_extensions/outerbounds/remote_config.pyi +1 -1
  86. metaflow-stubs/mf_extensions/outerbounds/toplevel/__init__.pyi +1 -1
  87. metaflow-stubs/mf_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.pyi +1 -1
  88. metaflow-stubs/multicore_utils.pyi +1 -1
  89. metaflow-stubs/ob_internal.pyi +1 -1
  90. metaflow-stubs/parameters.pyi +2 -2
  91. metaflow-stubs/plugins/__init__.pyi +9 -9
  92. metaflow-stubs/plugins/airflow/__init__.pyi +1 -1
  93. metaflow-stubs/plugins/airflow/airflow_utils.pyi +1 -1
  94. metaflow-stubs/plugins/airflow/exception.pyi +1 -1
  95. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +1 -1
  96. metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi +1 -1
  97. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +1 -1
  98. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +1 -1
  99. metaflow-stubs/plugins/argo/__init__.pyi +1 -1
  100. metaflow-stubs/plugins/argo/argo_client.pyi +1 -1
  101. metaflow-stubs/plugins/argo/argo_events.pyi +1 -1
  102. metaflow-stubs/plugins/argo/argo_workflows.pyi +2 -2
  103. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +2 -2
  104. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +2 -2
  105. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +2 -2
  106. metaflow-stubs/plugins/aws/__init__.pyi +1 -1
  107. metaflow-stubs/plugins/aws/aws_client.pyi +1 -1
  108. metaflow-stubs/plugins/aws/aws_utils.pyi +1 -1
  109. metaflow-stubs/plugins/aws/batch/__init__.pyi +1 -1
  110. metaflow-stubs/plugins/aws/batch/batch.pyi +1 -1
  111. metaflow-stubs/plugins/aws/batch/batch_client.pyi +1 -1
  112. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +1 -1
  113. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +1 -1
  114. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +3 -3
  115. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +1 -1
  116. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +1 -1
  117. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +1 -1
  118. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +1 -1
  119. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +1 -1
  120. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer.pyi +3 -3
  121. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer_objects.pyi +1 -1
  122. metaflow-stubs/plugins/azure/__init__.pyi +1 -1
  123. metaflow-stubs/plugins/azure/azure_credential.pyi +1 -1
  124. metaflow-stubs/plugins/azure/azure_exceptions.pyi +1 -1
  125. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +3 -3
  126. metaflow-stubs/plugins/azure/azure_utils.pyi +1 -1
  127. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +1 -1
  128. metaflow-stubs/plugins/azure/includefile_support.pyi +1 -1
  129. metaflow-stubs/plugins/cards/__init__.pyi +1 -1
  130. metaflow-stubs/plugins/cards/card_client.pyi +1 -1
  131. metaflow-stubs/plugins/cards/card_creator.pyi +1 -1
  132. metaflow-stubs/plugins/cards/card_datastore.pyi +1 -1
  133. metaflow-stubs/plugins/cards/card_decorator.pyi +1 -1
  134. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +1 -1
  135. metaflow-stubs/plugins/cards/card_modules/basic.pyi +2 -2
  136. metaflow-stubs/plugins/cards/card_modules/card.pyi +1 -1
  137. metaflow-stubs/plugins/cards/card_modules/components.pyi +2 -2
  138. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +1 -1
  139. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +1 -1
  140. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +1 -1
  141. metaflow-stubs/plugins/cards/card_resolver.pyi +1 -1
  142. metaflow-stubs/plugins/cards/component_serializer.pyi +1 -1
  143. metaflow-stubs/plugins/cards/exception.pyi +1 -1
  144. metaflow-stubs/plugins/catch_decorator.pyi +2 -2
  145. metaflow-stubs/plugins/datatools/__init__.pyi +1 -1
  146. metaflow-stubs/plugins/datatools/local.pyi +1 -1
  147. metaflow-stubs/plugins/datatools/s3/__init__.pyi +1 -1
  148. metaflow-stubs/plugins/datatools/s3/s3.pyi +2 -2
  149. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +1 -1
  150. metaflow-stubs/plugins/datatools/s3/s3util.pyi +1 -1
  151. metaflow-stubs/plugins/debug_logger.pyi +1 -1
  152. metaflow-stubs/plugins/debug_monitor.pyi +1 -1
  153. metaflow-stubs/plugins/environment_decorator.pyi +1 -1
  154. metaflow-stubs/plugins/events_decorator.pyi +1 -1
  155. metaflow-stubs/plugins/frameworks/__init__.pyi +1 -1
  156. metaflow-stubs/plugins/frameworks/pytorch.pyi +1 -1
  157. metaflow-stubs/plugins/gcp/__init__.pyi +1 -1
  158. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +3 -3
  159. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +1 -1
  160. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +1 -1
  161. metaflow-stubs/plugins/gcp/gs_utils.pyi +1 -1
  162. metaflow-stubs/plugins/gcp/includefile_support.pyi +1 -1
  163. metaflow-stubs/plugins/kubernetes/__init__.pyi +1 -1
  164. metaflow-stubs/plugins/kubernetes/kube_utils.pyi +1 -1
  165. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +1 -1
  166. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +1 -1
  167. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +1 -1
  168. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +1 -1
  169. metaflow-stubs/plugins/kubernetes/spot_monitor_sidecar.pyi +1 -1
  170. metaflow-stubs/plugins/ollama/__init__.pyi +2 -2
  171. metaflow-stubs/plugins/parallel_decorator.pyi +1 -1
  172. metaflow-stubs/plugins/perimeters.pyi +1 -1
  173. metaflow-stubs/plugins/project_decorator.pyi +1 -1
  174. metaflow-stubs/plugins/pypi/__init__.pyi +1 -1
  175. metaflow-stubs/plugins/pypi/conda_decorator.pyi +1 -1
  176. metaflow-stubs/plugins/pypi/conda_environment.pyi +4 -4
  177. metaflow-stubs/plugins/pypi/parsers.pyi +1 -1
  178. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +1 -1
  179. metaflow-stubs/plugins/pypi/pypi_environment.pyi +1 -1
  180. metaflow-stubs/plugins/pypi/utils.pyi +1 -1
  181. metaflow-stubs/plugins/resources_decorator.pyi +1 -1
  182. metaflow-stubs/plugins/retry_decorator.pyi +1 -1
  183. metaflow-stubs/plugins/secrets/__init__.pyi +1 -1
  184. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +2 -2
  185. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +1 -1
  186. metaflow-stubs/plugins/snowflake/__init__.pyi +1 -1
  187. metaflow-stubs/plugins/storage_executor.pyi +1 -1
  188. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +2 -2
  189. metaflow-stubs/plugins/timeout_decorator.pyi +2 -2
  190. metaflow-stubs/plugins/torchtune/__init__.pyi +3 -2
  191. metaflow-stubs/plugins/uv/__init__.pyi +1 -1
  192. metaflow-stubs/plugins/uv/uv_environment.pyi +1 -1
  193. metaflow-stubs/profilers/__init__.pyi +1 -1
  194. metaflow-stubs/pylint_wrapper.pyi +1 -1
  195. metaflow-stubs/runner/__init__.pyi +1 -1
  196. metaflow-stubs/runner/deployer.pyi +28 -28
  197. metaflow-stubs/runner/deployer_impl.pyi +2 -2
  198. metaflow-stubs/runner/metaflow_runner.pyi +2 -2
  199. metaflow-stubs/runner/nbdeploy.pyi +1 -1
  200. metaflow-stubs/runner/nbrun.pyi +1 -1
  201. metaflow-stubs/runner/subprocess_manager.pyi +1 -1
  202. metaflow-stubs/runner/utils.pyi +2 -2
  203. metaflow-stubs/system/__init__.pyi +1 -1
  204. metaflow-stubs/system/system_logger.pyi +2 -2
  205. metaflow-stubs/system/system_monitor.pyi +1 -1
  206. metaflow-stubs/tagging_util.pyi +1 -1
  207. metaflow-stubs/tuple_util.pyi +1 -1
  208. metaflow-stubs/user_configs/__init__.pyi +1 -1
  209. metaflow-stubs/user_configs/config_decorators.pyi +5 -5
  210. metaflow-stubs/user_configs/config_options.pyi +2 -2
  211. metaflow-stubs/user_configs/config_parameters.pyi +6 -6
  212. {ob_metaflow_stubs-6.0.3.184.dist-info → ob_metaflow_stubs-6.0.3.185.dist-info}/METADATA +1 -1
  213. ob_metaflow_stubs-6.0.3.185.dist-info/RECORD +216 -0
  214. ob_metaflow_stubs-6.0.3.184.dist-info/RECORD +0 -216
  215. {ob_metaflow_stubs-6.0.3.184.dist-info → ob_metaflow_stubs-6.0.3.185.dist-info}/WHEEL +0 -0
  216. {ob_metaflow_stubs-6.0.3.184.dist-info → ob_metaflow_stubs-6.0.3.185.dist-info}/top_level.txt +0 -0
@@ -1,7 +1,7 @@
1
1
  ######################################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
3
  # MF version: 2.15.18.1+obcheckpoint(0.2.1);ob(v1) #
4
- # Generated on 2025-06-25T20:19:31.361984 #
4
+ # Generated on 2025-06-26T22:38:03.208016 #
5
5
  ######################################################################################################
6
6
 
7
7
  from __future__ import annotations
@@ -35,18 +35,18 @@ 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
40
- from . import events as events
41
39
  from . import metaflow_git as metaflow_git
40
+ from . import tuple_util as tuple_util
41
+ from . import events as events
42
42
  from . import runner as runner
43
43
  from . import plugins as plugins
44
44
  from .mf_extensions.outerbounds.toplevel.global_aliases_for_metaflow_package import S3 as S3
45
45
  from . import includefile as includefile
46
46
  from .includefile import IncludeFile as IncludeFile
47
- from .plugins.pypi.parsers import conda_environment_yml_parser as conda_environment_yml_parser
48
47
  from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
49
48
  from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
49
+ from .plugins.pypi.parsers import conda_environment_yml_parser as conda_environment_yml_parser
50
50
  from . import client as client
51
51
  from .client.core import namespace as namespace
52
52
  from .client.core import get_namespace as get_namespace
@@ -155,7 +155,7 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
155
155
  """
156
156
  ...
157
157
 
158
- 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]]]:
158
+ def nvct(*, gpu: int, gpu_type: str) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
159
159
  """
160
160
  Specifies that this step should execute on DGX cloud.
161
161
 
@@ -166,125 +166,115 @@ def nvidia(*, gpu: int, gpu_type: str, queue_timeout: int) -> typing.Callable[[t
166
166
  Number of GPUs to use.
167
167
  gpu_type : str
168
168
  Type of Nvidia GPU to use.
169
- queue_timeout : int
170
- Time to keep the job in NVCF's queue.
171
169
  """
172
170
  ...
173
171
 
174
- @typing.overload
175
- 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]]]:
172
+ 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]]]:
176
173
  """
177
- Specifies a timeout for your step.
178
-
179
- This decorator is useful if this step may hang indefinitely.
180
-
181
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
182
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
183
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
184
-
185
- Note that all the values specified in parameters are added together so if you specify
186
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
174
+ Specifies that this step is used to deploy an instance of the app.
175
+ Requires that self.app_name, self.app_port, self.entrypoint and self.deployDir is set.
187
176
 
188
177
 
189
178
  Parameters
190
179
  ----------
191
- seconds : int, default 0
192
- Number of seconds to wait prior to timing out.
193
- minutes : int, default 0
194
- Number of minutes to wait prior to timing out.
195
- hours : int, default 0
196
- Number of hours to wait prior to timing out.
180
+ app_port : int
181
+ Number of GPUs to use.
182
+ app_name : str
183
+ Name of the app to deploy.
197
184
  """
198
185
  ...
199
186
 
200
- @typing.overload
201
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
202
- ...
203
-
204
- @typing.overload
205
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
206
- ...
207
-
208
- 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):
187
+ def ollama(*, models: list, backend: str, force_pull: bool, cache_update_policy: str, force_cache_update: bool, debug: bool, circuit_breaker_config: dict, timeout_config: dict) -> 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]]]:
209
188
  """
210
- Specifies a timeout for your step.
211
-
212
- This decorator is useful if this step may hang indefinitely.
213
-
214
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
215
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
216
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
189
+ This decorator is used to run Ollama APIs as Metaflow task sidecars.
217
190
 
218
- Note that all the values specified in parameters are added together so if you specify
219
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
191
+ User code call
192
+ --------------
193
+ @ollama(
194
+ models=[...],
195
+ ...
196
+ )
220
197
 
198
+ Valid backend options
199
+ ---------------------
200
+ - 'local': Run as a separate process on the local task machine.
201
+ - (TODO) 'managed': Outerbounds hosts and selects compute provider.
202
+ - (TODO) 'remote': Spin up separate instance to serve Ollama models.
221
203
 
222
- Parameters
223
- ----------
224
- seconds : int, default 0
225
- Number of seconds to wait prior to timing out.
226
- minutes : int, default 0
227
- Number of minutes to wait prior to timing out.
228
- hours : int, default 0
229
- Number of hours to wait prior to timing out.
230
- """
231
- ...
232
-
233
- def huggingface_hub(*, temp_dir_root: typing.Optional[str] = None, load: typing.Union[typing.List[str], typing.List[typing.Tuple[typing.Dict, str]], typing.List[typing.Tuple[str, str]], typing.List[typing.Dict], None]) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
234
- """
235
- Decorator that helps cache, version and store models/datasets from huggingface hub.
204
+ Valid model options
205
+ -------------------
206
+ Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
236
207
 
237
208
 
238
209
  Parameters
239
210
  ----------
240
- temp_dir_root : str, optional
241
- The root directory that will hold the temporary directory where objects will be downloaded.
242
-
243
- load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
244
- The list of repos (models/datasets) to load.
245
-
246
- Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
247
-
248
- - If repo (model/dataset) is not found in the datastore:
249
- - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
250
- - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
251
- - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
252
-
253
- - If repo is found in the datastore:
254
- - Loads it directly from datastore to local path (can be temporary directory or specified path)
211
+ models: list[str]
212
+ List of Ollama containers running models in sidecars.
213
+ backend: str
214
+ Determines where and how to run the Ollama process.
215
+ force_pull: bool
216
+ Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
217
+ cache_update_policy: str
218
+ Cache update policy: "auto", "force", or "never".
219
+ force_cache_update: bool
220
+ Simple override for "force" cache update policy.
221
+ debug: bool
222
+ Whether to turn on verbose debugging logs.
223
+ circuit_breaker_config: dict
224
+ Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
225
+ timeout_config: dict
226
+ Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
255
227
  """
256
228
  ...
257
229
 
258
230
  @typing.overload
259
- 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]]]:
231
+ def pypi(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
260
232
  """
261
- Specifies environment variables to be set prior to the execution of a step.
233
+ Specifies the PyPI packages for the step.
234
+
235
+ Information in this decorator will augment any
236
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
237
+ you can use `@pypi_base` to set packages required by all
238
+ steps and use `@pypi` to specify step-specific overrides.
262
239
 
263
240
 
264
241
  Parameters
265
242
  ----------
266
- vars : Dict[str, str], default {}
267
- Dictionary of environment variables to set.
243
+ packages : Dict[str, str], default: {}
244
+ Packages to use for this step. The key is the name of the package
245
+ and the value is the version to use.
246
+ python : str, optional, default: None
247
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
248
+ that the version used will correspond to the version of the Python interpreter used to start the run.
268
249
  """
269
250
  ...
270
251
 
271
252
  @typing.overload
272
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
253
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
273
254
  ...
274
255
 
275
256
  @typing.overload
276
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
257
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
277
258
  ...
278
259
 
279
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
260
+ def pypi(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
280
261
  """
281
- Specifies environment variables to be set prior to the execution of a step.
262
+ Specifies the PyPI packages for the step.
263
+
264
+ Information in this decorator will augment any
265
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
266
+ you can use `@pypi_base` to set packages required by all
267
+ steps and use `@pypi` to specify step-specific overrides.
282
268
 
283
269
 
284
270
  Parameters
285
271
  ----------
286
- vars : Dict[str, str], default {}
287
- Dictionary of environment variables to set.
272
+ packages : Dict[str, str], default: {}
273
+ Packages to use for this step. The key is the name of the package
274
+ and the value is the version to use.
275
+ python : str, optional, default: None
276
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
277
+ that the version used will correspond to the version of the Python interpreter used to start the run.
288
278
  """
289
279
  ...
290
280
 
@@ -397,139 +387,58 @@ def model(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
397
387
  ...
398
388
 
399
389
  @typing.overload
400
- 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]]]:
390
+ 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]]]:
401
391
  """
402
- Specifies the resources needed when executing this step.
403
-
404
- Use `@resources` to specify the resource requirements
405
- independently of the specific compute layer (`@batch`, `@kubernetes`).
406
-
407
- You can choose the compute layer on the command line by executing e.g.
408
- ```
409
- python myflow.py run --with batch
410
- ```
411
- or
412
- ```
413
- python myflow.py run --with kubernetes
414
- ```
415
- which executes the flow on the desired system using the
416
- requirements specified in `@resources`.
392
+ Specifies environment variables to be set prior to the execution of a step.
417
393
 
418
394
 
419
395
  Parameters
420
396
  ----------
421
- cpu : int, default 1
422
- Number of CPUs required for this step.
423
- gpu : int, optional, default None
424
- Number of GPUs required for this step.
425
- disk : int, optional, default None
426
- Disk size (in MB) required for this step. Only applies on Kubernetes.
427
- memory : int, default 4096
428
- Memory size (in MB) required for this step.
429
- shared_memory : int, optional, default None
430
- The value for the size (in MiB) of the /dev/shm volume for this step.
431
- This parameter maps to the `--shm-size` option in Docker.
397
+ vars : Dict[str, str], default {}
398
+ Dictionary of environment variables to set.
432
399
  """
433
400
  ...
434
401
 
435
402
  @typing.overload
436
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
403
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
437
404
  ...
438
405
 
439
406
  @typing.overload
440
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
407
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
441
408
  ...
442
409
 
443
- 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):
410
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
444
411
  """
445
- Specifies the resources needed when executing this step.
446
-
447
- Use `@resources` to specify the resource requirements
448
- independently of the specific compute layer (`@batch`, `@kubernetes`).
449
-
450
- You can choose the compute layer on the command line by executing e.g.
451
- ```
452
- python myflow.py run --with batch
453
- ```
454
- or
455
- ```
456
- python myflow.py run --with kubernetes
457
- ```
458
- which executes the flow on the desired system using the
459
- requirements specified in `@resources`.
412
+ Specifies environment variables to be set prior to the execution of a step.
460
413
 
461
414
 
462
415
  Parameters
463
416
  ----------
464
- cpu : int, default 1
465
- Number of CPUs required for this step.
466
- gpu : int, optional, default None
467
- Number of GPUs required for this step.
468
- disk : int, optional, default None
469
- Disk size (in MB) required for this step. Only applies on Kubernetes.
470
- memory : int, default 4096
471
- Memory size (in MB) required for this step.
472
- shared_memory : int, optional, default None
473
- The value for the size (in MiB) of the /dev/shm volume for this step.
474
- This parameter maps to the `--shm-size` option in Docker.
417
+ vars : Dict[str, str], default {}
418
+ Dictionary of environment variables to set.
475
419
  """
476
420
  ...
477
421
 
478
- @typing.overload
479
- def pypi(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
422
+ 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]]]:
480
423
  """
481
- Specifies the PyPI packages for the step.
482
-
483
- Information in this decorator will augment any
484
- attributes set in the `@pyi_base` flow-level decorator. Hence,
485
- you can use `@pypi_base` to set packages required by all
486
- steps and use `@pypi` to specify step-specific overrides.
424
+ Specifies that this step should execute on DGX cloud.
487
425
 
488
426
 
489
427
  Parameters
490
428
  ----------
491
- packages : Dict[str, str], default: {}
492
- Packages to use for this step. The key is the name of the package
493
- and the value is the version to use.
494
- python : str, optional, default: None
495
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
496
- that the version used will correspond to the version of the Python interpreter used to start the run.
429
+ gpu : int
430
+ Number of GPUs to use.
431
+ gpu_type : str
432
+ Type of Nvidia GPU to use.
433
+ queue_timeout : int
434
+ Time to keep the job in NVCF's queue.
497
435
  """
498
436
  ...
499
437
 
500
438
  @typing.overload
501
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
502
- ...
503
-
504
- @typing.overload
505
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
506
- ...
507
-
508
- def pypi(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
439
+ def checkpoint(*, load_policy: str = 'fresh', temp_dir_root: str = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
509
440
  """
510
- Specifies the PyPI packages for the step.
511
-
512
- Information in this decorator will augment any
513
- attributes set in the `@pyi_base` flow-level decorator. Hence,
514
- you can use `@pypi_base` to set packages required by all
515
- steps and use `@pypi` to specify step-specific overrides.
516
-
517
-
518
- Parameters
519
- ----------
520
- packages : Dict[str, str], default: {}
521
- Packages to use for this step. The key is the name of the package
522
- and the value is the version to use.
523
- python : str, optional, default: None
524
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
525
- that the version used will correspond to the version of the Python interpreter used to start the run.
526
- """
527
- ...
528
-
529
- @typing.overload
530
- def checkpoint(*, load_policy: str = 'fresh', temp_dir_root: str = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
531
- """
532
- Enables checkpointing for a step.
441
+ Enables checkpointing for a step.
533
442
 
534
443
 
535
444
 
@@ -583,34 +492,22 @@ def checkpoint(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None
583
492
  """
584
493
  ...
585
494
 
586
- def nvct(*, gpu: int, gpu_type: str) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
587
- """
588
- Specifies that this step should execute on DGX cloud.
589
-
590
-
591
- Parameters
592
- ----------
593
- gpu : int
594
- Number of GPUs to use.
595
- gpu_type : str
596
- Type of Nvidia GPU to use.
597
- """
598
- ...
599
-
600
495
  @typing.overload
601
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
496
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
602
497
  """
603
- Internal decorator to support Fast bakery
498
+ Decorator prototype for all step decorators. This function gets specialized
499
+ and imported for all decorators types by _import_plugin_decorators().
604
500
  """
605
501
  ...
606
502
 
607
503
  @typing.overload
608
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
504
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
609
505
  ...
610
506
 
611
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
507
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
612
508
  """
613
- Internal decorator to support Fast bakery
509
+ Decorator prototype for all step decorators. This function gets specialized
510
+ and imported for all decorators types by _import_plugin_decorators().
614
511
  """
615
512
  ...
616
513
 
@@ -674,72 +571,141 @@ def conda(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
674
571
  ...
675
572
 
676
573
  @typing.overload
677
- def retry(*, times: int = 3, minutes_between_retries: int = 2) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
574
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
678
575
  """
679
- Specifies the number of times the task corresponding
680
- to a step needs to be retried.
576
+ Internal decorator to support Fast bakery
577
+ """
578
+ ...
579
+
580
+ @typing.overload
581
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
582
+ ...
583
+
584
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
585
+ """
586
+ Internal decorator to support Fast bakery
587
+ """
588
+ ...
589
+
590
+ def vllm(*, model: str, backend: str, debug: bool, kwargs: 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]]]:
591
+ """
592
+ This decorator is used to run vllm APIs as Metaflow task sidecars.
681
593
 
682
- This decorator is useful for handling transient errors, such as networking issues.
683
- If your task contains operations that can't be retried safely, e.g. database updates,
684
- it is advisable to annotate it with `@retry(times=0)`.
594
+ User code call
595
+ --------------
596
+ @vllm(
597
+ model="...",
598
+ ...
599
+ )
685
600
 
686
- This can be used in conjunction with the `@catch` decorator. The `@catch`
687
- decorator will execute a no-op task after all retries have been exhausted,
688
- ensuring that the flow execution can continue.
601
+ Valid backend options
602
+ ---------------------
603
+ - 'local': Run as a separate process on the local task machine.
604
+
605
+ Valid model options
606
+ -------------------
607
+ Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
608
+
609
+ NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
610
+ If you need multiple models, you must create multiple @vllm decorators.
689
611
 
690
612
 
691
613
  Parameters
692
614
  ----------
693
- times : int, default 3
694
- Number of times to retry this task.
695
- minutes_between_retries : int, default 2
696
- Number of minutes between retries.
615
+ model: str
616
+ HuggingFace model identifier to be served by vLLM.
617
+ backend: str
618
+ Determines where and how to run the vLLM process.
619
+ debug: bool
620
+ Whether to turn on verbose debugging logs.
621
+ kwargs : Any
622
+ Any other keyword arguments are passed directly to the vLLM engine.
623
+ This allows for flexible configuration of vLLM server settings.
624
+ For example, `tensor_parallel_size=2`.
697
625
  """
698
626
  ...
699
627
 
700
628
  @typing.overload
701
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
629
+ 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]]]:
630
+ """
631
+ Specifies a timeout for your step.
632
+
633
+ This decorator is useful if this step may hang indefinitely.
634
+
635
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
636
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
637
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
638
+
639
+ Note that all the values specified in parameters are added together so if you specify
640
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
641
+
642
+
643
+ Parameters
644
+ ----------
645
+ seconds : int, default 0
646
+ Number of seconds to wait prior to timing out.
647
+ minutes : int, default 0
648
+ Number of minutes to wait prior to timing out.
649
+ hours : int, default 0
650
+ Number of hours to wait prior to timing out.
651
+ """
702
652
  ...
703
653
 
704
654
  @typing.overload
705
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
655
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
706
656
  ...
707
657
 
708
- def retry(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, times: int = 3, minutes_between_retries: int = 2):
658
+ @typing.overload
659
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
660
+ ...
661
+
662
+ 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):
709
663
  """
710
- Specifies the number of times the task corresponding
711
- to a step needs to be retried.
664
+ Specifies a timeout for your step.
712
665
 
713
- This decorator is useful for handling transient errors, such as networking issues.
714
- If your task contains operations that can't be retried safely, e.g. database updates,
715
- it is advisable to annotate it with `@retry(times=0)`.
666
+ This decorator is useful if this step may hang indefinitely.
716
667
 
717
- This can be used in conjunction with the `@catch` decorator. The `@catch`
718
- decorator will execute a no-op task after all retries have been exhausted,
719
- ensuring that the flow execution can continue.
668
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
669
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
670
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
671
+
672
+ Note that all the values specified in parameters are added together so if you specify
673
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
720
674
 
721
675
 
722
676
  Parameters
723
677
  ----------
724
- times : int, default 3
725
- Number of times to retry this task.
726
- minutes_between_retries : int, default 2
727
- Number of minutes between retries.
678
+ seconds : int, default 0
679
+ Number of seconds to wait prior to timing out.
680
+ minutes : int, default 0
681
+ Number of minutes to wait prior to timing out.
682
+ hours : int, default 0
683
+ Number of hours to wait prior to timing out.
728
684
  """
729
685
  ...
730
686
 
731
- 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]]]:
687
+ def huggingface_hub(*, temp_dir_root: typing.Optional[str] = None, load: typing.Union[typing.List[str], typing.List[typing.Tuple[typing.Dict, str]], typing.List[typing.Tuple[str, str]], typing.List[typing.Dict], None]) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
732
688
  """
733
- Specifies that this step is used to deploy an instance of the app.
734
- Requires that self.app_name, self.app_port, self.entrypoint and self.deployDir is set.
689
+ Decorator that helps cache, version and store models/datasets from huggingface hub.
735
690
 
736
691
 
737
692
  Parameters
738
693
  ----------
739
- app_port : int
740
- Number of GPUs to use.
741
- app_name : str
742
- Name of the app to deploy.
694
+ temp_dir_root : str, optional
695
+ The root directory that will hold the temporary directory where objects will be downloaded.
696
+
697
+ load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
698
+ The list of repos (models/datasets) to load.
699
+
700
+ Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
701
+
702
+ - If repo (model/dataset) is not found in the datastore:
703
+ - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
704
+ - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
705
+ - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
706
+
707
+ - If repo is found in the datastore:
708
+ - Loads it directly from datastore to local path (can be temporary directory or specified path)
743
709
  """
744
710
  ...
745
711
 
@@ -778,44 +744,6 @@ def secrets(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None],
778
744
  """
779
745
  ...
780
746
 
781
- def vllm(*, model: str, backend: str, debug: bool, kwargs: 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]]]:
782
- """
783
- This decorator is used to run vllm APIs as Metaflow task sidecars.
784
-
785
- User code call
786
- --------------
787
- @vllm(
788
- model="...",
789
- ...
790
- )
791
-
792
- Valid backend options
793
- ---------------------
794
- - 'local': Run as a separate process on the local task machine.
795
-
796
- Valid model options
797
- -------------------
798
- Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
799
-
800
- NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
801
- If you need multiple models, you must create multiple @vllm decorators.
802
-
803
-
804
- Parameters
805
- ----------
806
- model: str
807
- HuggingFace model identifier to be served by vLLM.
808
- backend: str
809
- Determines where and how to run the vLLM process.
810
- debug: bool
811
- Whether to turn on verbose debugging logs.
812
- kwargs : Any
813
- Any other keyword arguments are passed directly to the vLLM engine.
814
- This allows for flexible configuration of vLLM server settings.
815
- For example, `tensor_parallel_size=2`.
816
- """
817
- ...
818
-
819
747
  def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: typing.Optional[str] = None, image_pull_policy: str = 'KUBERNETES_IMAGE_PULL_POLICY', image_pull_secrets: typing.List[str] = [], 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', security_context: typing.Optional[typing.Dict[str, typing.Any]] = 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]]]:
820
748
  """
821
749
  Specifies that this step should execute on Kubernetes.
@@ -905,46 +833,58 @@ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: ty
905
833
  """
906
834
  ...
907
835
 
908
- def ollama(*, models: list, backend: str, force_pull: bool, cache_update_policy: str, force_cache_update: bool, debug: bool, circuit_breaker_config: dict, timeout_config: dict) -> 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]]]:
836
+ @typing.overload
837
+ def retry(*, times: int = 3, minutes_between_retries: int = 2) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
909
838
  """
910
- This decorator is used to run Ollama APIs as Metaflow task sidecars.
839
+ Specifies the number of times the task corresponding
840
+ to a step needs to be retried.
911
841
 
912
- User code call
913
- --------------
914
- @ollama(
915
- models=[...],
916
- ...
917
- )
842
+ This decorator is useful for handling transient errors, such as networking issues.
843
+ If your task contains operations that can't be retried safely, e.g. database updates,
844
+ it is advisable to annotate it with `@retry(times=0)`.
918
845
 
919
- Valid backend options
920
- ---------------------
921
- - 'local': Run as a separate process on the local task machine.
922
- - (TODO) 'managed': Outerbounds hosts and selects compute provider.
923
- - (TODO) 'remote': Spin up separate instance to serve Ollama models.
846
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
847
+ decorator will execute a no-op task after all retries have been exhausted,
848
+ ensuring that the flow execution can continue.
924
849
 
925
- Valid model options
926
- -------------------
927
- Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
850
+
851
+ Parameters
852
+ ----------
853
+ times : int, default 3
854
+ Number of times to retry this task.
855
+ minutes_between_retries : int, default 2
856
+ Number of minutes between retries.
857
+ """
858
+ ...
859
+
860
+ @typing.overload
861
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
862
+ ...
863
+
864
+ @typing.overload
865
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
866
+ ...
867
+
868
+ def retry(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, times: int = 3, minutes_between_retries: int = 2):
869
+ """
870
+ Specifies the number of times the task corresponding
871
+ to a step needs to be retried.
872
+
873
+ This decorator is useful for handling transient errors, such as networking issues.
874
+ If your task contains operations that can't be retried safely, e.g. database updates,
875
+ it is advisable to annotate it with `@retry(times=0)`.
876
+
877
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
878
+ decorator will execute a no-op task after all retries have been exhausted,
879
+ ensuring that the flow execution can continue.
928
880
 
929
881
 
930
882
  Parameters
931
883
  ----------
932
- models: list[str]
933
- List of Ollama containers running models in sidecars.
934
- backend: str
935
- Determines where and how to run the Ollama process.
936
- force_pull: bool
937
- Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
938
- cache_update_policy: str
939
- Cache update policy: "auto", "force", or "never".
940
- force_cache_update: bool
941
- Simple override for "force" cache update policy.
942
- debug: bool
943
- Whether to turn on verbose debugging logs.
944
- circuit_breaker_config: dict
945
- Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
946
- timeout_config: dict
947
- Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
884
+ times : int, default 3
885
+ Number of times to retry this task.
886
+ minutes_between_retries : int, default 2
887
+ Number of minutes between retries.
948
888
  """
949
889
  ...
950
890
 
@@ -998,72 +938,116 @@ def card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typ
998
938
  ...
999
939
 
1000
940
  @typing.overload
1001
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
941
+ 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]]]:
1002
942
  """
1003
- Decorator prototype for all step decorators. This function gets specialized
1004
- and imported for all decorators types by _import_plugin_decorators().
943
+ Specifies the resources needed when executing this step.
944
+
945
+ Use `@resources` to specify the resource requirements
946
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
947
+
948
+ You can choose the compute layer on the command line by executing e.g.
949
+ ```
950
+ python myflow.py run --with batch
951
+ ```
952
+ or
953
+ ```
954
+ python myflow.py run --with kubernetes
955
+ ```
956
+ which executes the flow on the desired system using the
957
+ requirements specified in `@resources`.
958
+
959
+
960
+ Parameters
961
+ ----------
962
+ cpu : int, default 1
963
+ Number of CPUs required for this step.
964
+ gpu : int, optional, default None
965
+ Number of GPUs required for this step.
966
+ disk : int, optional, default None
967
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
968
+ memory : int, default 4096
969
+ Memory size (in MB) required for this step.
970
+ shared_memory : int, optional, default None
971
+ The value for the size (in MiB) of the /dev/shm volume for this step.
972
+ This parameter maps to the `--shm-size` option in Docker.
1005
973
  """
1006
974
  ...
1007
975
 
1008
976
  @typing.overload
1009
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
977
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1010
978
  ...
1011
979
 
1012
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1013
- """
1014
- Decorator prototype for all step decorators. This function gets specialized
1015
- and imported for all decorators types by _import_plugin_decorators().
1016
- """
980
+ @typing.overload
981
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1017
982
  ...
1018
983
 
1019
- @typing.overload
1020
- 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]]:
984
+ 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):
1021
985
  """
1022
- Specifies the Conda environment for all steps of the flow.
986
+ Specifies the resources needed when executing this step.
1023
987
 
1024
- Use `@conda_base` to set common libraries required by all
1025
- steps and use `@conda` to specify step-specific additions.
988
+ Use `@resources` to specify the resource requirements
989
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
990
+
991
+ You can choose the compute layer on the command line by executing e.g.
992
+ ```
993
+ python myflow.py run --with batch
994
+ ```
995
+ or
996
+ ```
997
+ python myflow.py run --with kubernetes
998
+ ```
999
+ which executes the flow on the desired system using the
1000
+ requirements specified in `@resources`.
1026
1001
 
1027
1002
 
1028
1003
  Parameters
1029
1004
  ----------
1030
- packages : Dict[str, str], default {}
1031
- Packages to use for this flow. The key is the name of the package
1032
- and the value is the version to use.
1033
- libraries : Dict[str, str], default {}
1034
- Supported for backward compatibility. When used with packages, packages will take precedence.
1035
- python : str, optional, default None
1036
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1037
- that the version used will correspond to the version of the Python interpreter used to start the run.
1038
- disabled : bool, default False
1039
- If set to True, disables Conda.
1005
+ cpu : int, default 1
1006
+ Number of CPUs required for this step.
1007
+ gpu : int, optional, default None
1008
+ Number of GPUs required for this step.
1009
+ disk : int, optional, default None
1010
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
1011
+ memory : int, default 4096
1012
+ Memory size (in MB) required for this step.
1013
+ shared_memory : int, optional, default None
1014
+ The value for the size (in MiB) of the /dev/shm volume for this step.
1015
+ This parameter maps to the `--shm-size` option in Docker.
1040
1016
  """
1041
1017
  ...
1042
1018
 
1043
- @typing.overload
1044
- def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1045
- ...
1046
-
1047
- 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):
1019
+ def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1048
1020
  """
1049
- Specifies the Conda environment for all steps of the flow.
1021
+ Specifies what flows belong to the same project.
1050
1022
 
1051
- Use `@conda_base` to set common libraries required by all
1052
- steps and use `@conda` to specify step-specific additions.
1023
+ A project-specific namespace is created for all flows that
1024
+ use the same `@project(name)`.
1053
1025
 
1054
1026
 
1055
1027
  Parameters
1056
1028
  ----------
1057
- packages : Dict[str, str], default {}
1058
- Packages to use for this flow. The key is the name of the package
1059
- and the value is the version to use.
1060
- libraries : Dict[str, str], default {}
1061
- Supported for backward compatibility. When used with packages, packages will take precedence.
1062
- python : str, optional, default None
1063
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1064
- that the version used will correspond to the version of the Python interpreter used to start the run.
1065
- disabled : bool, default False
1066
- If set to True, disables Conda.
1029
+ name : str
1030
+ Project name. Make sure that the name is unique amongst all
1031
+ projects that use the same production scheduler. The name may
1032
+ contain only lowercase alphanumeric characters and underscores.
1033
+
1034
+ branch : Optional[str], default None
1035
+ The branch to use. If not specified, the branch is set to
1036
+ `user.<username>` unless `production` is set to `True`. This can
1037
+ also be set on the command line using `--branch` as a top-level option.
1038
+ It is an error to specify `branch` in the decorator and on the command line.
1039
+
1040
+ production : bool, default False
1041
+ Whether or not the branch is the production branch. This can also be set on the
1042
+ command line using `--production` as a top-level option. It is an error to specify
1043
+ `production` in the decorator and on the command line.
1044
+ The project branch name will be:
1045
+ - if `branch` is specified:
1046
+ - if `production` is True: `prod.<branch>`
1047
+ - if `production` is False: `test.<branch>`
1048
+ - if `branch` is not specified:
1049
+ - if `production` is True: `prod`
1050
+ - if `production` is False: `user.<username>`
1067
1051
  """
1068
1052
  ...
1069
1053
 
@@ -1224,6 +1208,47 @@ def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None)
1224
1208
  """
1225
1209
  ...
1226
1210
 
1211
+ @typing.overload
1212
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1213
+ """
1214
+ Specifies the PyPI packages for all steps of the flow.
1215
+
1216
+ Use `@pypi_base` to set common packages required by all
1217
+ steps and use `@pypi` to specify step-specific overrides.
1218
+
1219
+ Parameters
1220
+ ----------
1221
+ packages : Dict[str, str], default: {}
1222
+ Packages to use for this flow. The key is the name of the package
1223
+ and the value is the version to use.
1224
+ python : str, optional, default: None
1225
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1226
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1227
+ """
1228
+ ...
1229
+
1230
+ @typing.overload
1231
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1232
+ ...
1233
+
1234
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1235
+ """
1236
+ Specifies the PyPI packages for all steps of the flow.
1237
+
1238
+ Use `@pypi_base` to set common packages required by all
1239
+ steps and use `@pypi` to specify step-specific overrides.
1240
+
1241
+ Parameters
1242
+ ----------
1243
+ packages : Dict[str, str], default: {}
1244
+ Packages to use for this flow. The key is the name of the package
1245
+ and the value is the version to use.
1246
+ python : str, optional, default: None
1247
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1248
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1249
+ """
1250
+ ...
1251
+
1227
1252
  @typing.overload
1228
1253
  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]]:
1229
1254
  """
@@ -1275,6 +1300,100 @@ def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly:
1275
1300
  """
1276
1301
  ...
1277
1302
 
1303
+ @typing.overload
1304
+ 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]]:
1305
+ """
1306
+ Specifies the Conda environment for all steps of the flow.
1307
+
1308
+ Use `@conda_base` to set common libraries required by all
1309
+ steps and use `@conda` to specify step-specific additions.
1310
+
1311
+
1312
+ Parameters
1313
+ ----------
1314
+ packages : Dict[str, str], default {}
1315
+ Packages to use for this flow. The key is the name of the package
1316
+ and the value is the version to use.
1317
+ libraries : Dict[str, str], default {}
1318
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1319
+ python : str, optional, default None
1320
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1321
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1322
+ disabled : bool, default False
1323
+ If set to True, disables Conda.
1324
+ """
1325
+ ...
1326
+
1327
+ @typing.overload
1328
+ def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1329
+ ...
1330
+
1331
+ 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):
1332
+ """
1333
+ Specifies the Conda environment for all steps of the flow.
1334
+
1335
+ Use `@conda_base` to set common libraries required by all
1336
+ steps and use `@conda` to specify step-specific additions.
1337
+
1338
+
1339
+ Parameters
1340
+ ----------
1341
+ packages : Dict[str, str], default {}
1342
+ Packages to use for this flow. The key is the name of the package
1343
+ and the value is the version to use.
1344
+ libraries : Dict[str, str], default {}
1345
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1346
+ python : str, optional, default None
1347
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1348
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1349
+ disabled : bool, default False
1350
+ If set to True, disables Conda.
1351
+ """
1352
+ ...
1353
+
1354
+ 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]]:
1355
+ """
1356
+ 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)
1357
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1358
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1359
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
1360
+ starts only after all sensors finish.
1361
+
1362
+
1363
+ Parameters
1364
+ ----------
1365
+ timeout : int
1366
+ Time, in seconds before the task times out and fails. (Default: 3600)
1367
+ poke_interval : int
1368
+ Time in seconds that the job should wait in between each try. (Default: 60)
1369
+ mode : str
1370
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1371
+ exponential_backoff : bool
1372
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1373
+ pool : str
1374
+ the slot pool this task should run in,
1375
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1376
+ soft_fail : bool
1377
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1378
+ name : str
1379
+ Name of the sensor on Airflow
1380
+ description : str
1381
+ Description of sensor in the Airflow UI
1382
+ bucket_key : Union[str, List[str]]
1383
+ The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1384
+ When it's specified as a full s3:// url, please leave `bucket_name` as None
1385
+ bucket_name : str
1386
+ Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1387
+ When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1388
+ wildcard_match : bool
1389
+ whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1390
+ aws_conn_id : str
1391
+ a reference to the s3 connection on Airflow. (Default: None)
1392
+ verify : bool
1393
+ Whether or not to verify SSL certificates for S3 connection. (Default: None)
1394
+ """
1395
+ ...
1396
+
1278
1397
  @typing.overload
1279
1398
  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]]:
1280
1399
  """
@@ -1376,125 +1495,6 @@ def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *
1376
1495
  """
1377
1496
  ...
1378
1497
 
1379
- 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]]:
1380
- """
1381
- 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)
1382
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1383
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1384
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
1385
- starts only after all sensors finish.
1386
-
1387
-
1388
- Parameters
1389
- ----------
1390
- timeout : int
1391
- Time, in seconds before the task times out and fails. (Default: 3600)
1392
- poke_interval : int
1393
- Time in seconds that the job should wait in between each try. (Default: 60)
1394
- mode : str
1395
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1396
- exponential_backoff : bool
1397
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1398
- pool : str
1399
- the slot pool this task should run in,
1400
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1401
- soft_fail : bool
1402
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1403
- name : str
1404
- Name of the sensor on Airflow
1405
- description : str
1406
- Description of sensor in the Airflow UI
1407
- bucket_key : Union[str, List[str]]
1408
- The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1409
- When it's specified as a full s3:// url, please leave `bucket_name` as None
1410
- bucket_name : str
1411
- Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1412
- When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1413
- wildcard_match : bool
1414
- whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1415
- aws_conn_id : str
1416
- a reference to the s3 connection on Airflow. (Default: None)
1417
- verify : bool
1418
- Whether or not to verify SSL certificates for S3 connection. (Default: None)
1419
- """
1420
- ...
1421
-
1422
- def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1423
- """
1424
- Specifies what flows belong to the same project.
1425
-
1426
- A project-specific namespace is created for all flows that
1427
- use the same `@project(name)`.
1428
-
1429
-
1430
- Parameters
1431
- ----------
1432
- name : str
1433
- Project name. Make sure that the name is unique amongst all
1434
- projects that use the same production scheduler. The name may
1435
- contain only lowercase alphanumeric characters and underscores.
1436
-
1437
- branch : Optional[str], default None
1438
- The branch to use. If not specified, the branch is set to
1439
- `user.<username>` unless `production` is set to `True`. This can
1440
- also be set on the command line using `--branch` as a top-level option.
1441
- It is an error to specify `branch` in the decorator and on the command line.
1442
-
1443
- production : bool, default False
1444
- Whether or not the branch is the production branch. This can also be set on the
1445
- command line using `--production` as a top-level option. It is an error to specify
1446
- `production` in the decorator and on the command line.
1447
- The project branch name will be:
1448
- - if `branch` is specified:
1449
- - if `production` is True: `prod.<branch>`
1450
- - if `production` is False: `test.<branch>`
1451
- - if `branch` is not specified:
1452
- - if `production` is True: `prod`
1453
- - if `production` is False: `user.<username>`
1454
- """
1455
- ...
1456
-
1457
- @typing.overload
1458
- def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1459
- """
1460
- Specifies the PyPI packages for all steps of the flow.
1461
-
1462
- Use `@pypi_base` to set common packages required by all
1463
- steps and use `@pypi` to specify step-specific overrides.
1464
-
1465
- Parameters
1466
- ----------
1467
- packages : Dict[str, str], default: {}
1468
- Packages to use for this flow. The key is the name of the package
1469
- and the value is the version to use.
1470
- python : str, optional, default: None
1471
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1472
- that the version used will correspond to the version of the Python interpreter used to start the run.
1473
- """
1474
- ...
1475
-
1476
- @typing.overload
1477
- def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1478
- ...
1479
-
1480
- def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1481
- """
1482
- Specifies the PyPI packages for all steps of the flow.
1483
-
1484
- Use `@pypi_base` to set common packages required by all
1485
- steps and use `@pypi` to specify step-specific overrides.
1486
-
1487
- Parameters
1488
- ----------
1489
- packages : Dict[str, str], default: {}
1490
- Packages to use for this flow. The key is the name of the package
1491
- and the value is the version to use.
1492
- python : str, optional, default: None
1493
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1494
- that the version used will correspond to the version of the Python interpreter used to start the run.
1495
- """
1496
- ...
1497
-
1498
1498
  @typing.overload
1499
1499
  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]]:
1500
1500
  """