ob-metaflow-stubs 6.0.3.181__py2.py3-none-any.whl → 6.0.3.182rc1__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 (220) hide show
  1. metaflow-stubs/__init__.pyi +753 -752
  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 +5 -5
  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 +4 -4
  12. metaflow-stubs/generated_for.txt +1 -1
  13. metaflow-stubs/includefile.pyi +1 -1
  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 +139 -139
  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 +2 -2
  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 +2 -2
  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 +2 -2
  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 +2 -2
  43. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/core.pyi +2 -2
  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 +1 -1
  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 +2 -2
  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 +2 -2
  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/fast_bakery/__init__.pyi +6 -0
  75. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/baker.pyi +51 -0
  76. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/docker_environment.pyi +65 -0
  77. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/fast_bakery.pyi +74 -0
  78. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/__init__.pyi +1 -1
  79. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/pod_killer.pyi +1 -1
  80. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/__init__.pyi +1 -1
  81. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/constants.pyi +1 -1
  82. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/exceptions.pyi +1 -1
  83. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/ollama.pyi +1 -1
  84. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/status_card.pyi +1 -1
  85. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/__init__.pyi +1 -1
  86. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/snowflake.pyi +1 -1
  87. metaflow-stubs/mf_extensions/outerbounds/profilers/__init__.pyi +1 -1
  88. metaflow-stubs/mf_extensions/outerbounds/profilers/gpu.pyi +1 -1
  89. metaflow-stubs/mf_extensions/outerbounds/remote_config.pyi +1 -1
  90. metaflow-stubs/mf_extensions/outerbounds/toplevel/__init__.pyi +1 -1
  91. metaflow-stubs/mf_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.pyi +2 -1
  92. metaflow-stubs/multicore_utils.pyi +1 -1
  93. metaflow-stubs/ob_internal.pyi +12 -0
  94. metaflow-stubs/parameters.pyi +2 -2
  95. metaflow-stubs/plugins/__init__.pyi +10 -10
  96. metaflow-stubs/plugins/airflow/__init__.pyi +1 -1
  97. metaflow-stubs/plugins/airflow/airflow_utils.pyi +1 -1
  98. metaflow-stubs/plugins/airflow/exception.pyi +1 -1
  99. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +1 -1
  100. metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi +1 -1
  101. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +1 -1
  102. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +1 -1
  103. metaflow-stubs/plugins/argo/__init__.pyi +1 -1
  104. metaflow-stubs/plugins/argo/argo_client.pyi +1 -1
  105. metaflow-stubs/plugins/argo/argo_events.pyi +1 -1
  106. metaflow-stubs/plugins/argo/argo_workflows.pyi +2 -2
  107. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +2 -2
  108. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +1 -1
  109. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +1 -1
  110. metaflow-stubs/plugins/aws/__init__.pyi +1 -1
  111. metaflow-stubs/plugins/aws/aws_client.pyi +1 -1
  112. metaflow-stubs/plugins/aws/aws_utils.pyi +1 -1
  113. metaflow-stubs/plugins/aws/batch/__init__.pyi +1 -1
  114. metaflow-stubs/plugins/aws/batch/batch.pyi +1 -1
  115. metaflow-stubs/plugins/aws/batch/batch_client.pyi +1 -1
  116. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +1 -1
  117. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +1 -1
  118. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +3 -3
  119. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +1 -1
  120. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +1 -1
  121. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +1 -1
  122. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +1 -1
  123. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +1 -1
  124. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer.pyi +1 -1
  125. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer_objects.pyi +1 -1
  126. metaflow-stubs/plugins/azure/__init__.pyi +1 -1
  127. metaflow-stubs/plugins/azure/azure_credential.pyi +1 -1
  128. metaflow-stubs/plugins/azure/azure_exceptions.pyi +1 -1
  129. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +3 -3
  130. metaflow-stubs/plugins/azure/azure_utils.pyi +1 -1
  131. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +1 -1
  132. metaflow-stubs/plugins/azure/includefile_support.pyi +1 -1
  133. metaflow-stubs/plugins/cards/__init__.pyi +5 -5
  134. metaflow-stubs/plugins/cards/card_client.pyi +1 -1
  135. metaflow-stubs/plugins/cards/card_creator.pyi +1 -1
  136. metaflow-stubs/plugins/cards/card_datastore.pyi +1 -1
  137. metaflow-stubs/plugins/cards/card_decorator.pyi +1 -1
  138. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +1 -1
  139. metaflow-stubs/plugins/cards/card_modules/basic.pyi +1 -1
  140. metaflow-stubs/plugins/cards/card_modules/card.pyi +1 -1
  141. metaflow-stubs/plugins/cards/card_modules/components.pyi +2 -2
  142. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +1 -1
  143. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +1 -1
  144. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +1 -1
  145. metaflow-stubs/plugins/cards/card_resolver.pyi +1 -1
  146. metaflow-stubs/plugins/cards/component_serializer.pyi +1 -1
  147. metaflow-stubs/plugins/cards/exception.pyi +1 -1
  148. metaflow-stubs/plugins/catch_decorator.pyi +2 -2
  149. metaflow-stubs/plugins/datatools/__init__.pyi +1 -1
  150. metaflow-stubs/plugins/datatools/local.pyi +1 -1
  151. metaflow-stubs/plugins/datatools/s3/__init__.pyi +1 -1
  152. metaflow-stubs/plugins/datatools/s3/s3.pyi +1 -1
  153. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +1 -1
  154. metaflow-stubs/plugins/datatools/s3/s3util.pyi +1 -1
  155. metaflow-stubs/plugins/debug_logger.pyi +1 -1
  156. metaflow-stubs/plugins/debug_monitor.pyi +1 -1
  157. metaflow-stubs/plugins/environment_decorator.pyi +1 -1
  158. metaflow-stubs/plugins/events_decorator.pyi +1 -1
  159. metaflow-stubs/plugins/frameworks/__init__.pyi +1 -1
  160. metaflow-stubs/plugins/frameworks/pytorch.pyi +1 -1
  161. metaflow-stubs/plugins/gcp/__init__.pyi +1 -1
  162. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +3 -3
  163. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +1 -1
  164. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +1 -1
  165. metaflow-stubs/plugins/gcp/gs_utils.pyi +1 -1
  166. metaflow-stubs/plugins/gcp/includefile_support.pyi +1 -1
  167. metaflow-stubs/plugins/kubernetes/__init__.pyi +1 -1
  168. metaflow-stubs/plugins/kubernetes/kube_utils.pyi +1 -1
  169. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +1 -1
  170. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +1 -1
  171. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +1 -1
  172. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +1 -1
  173. metaflow-stubs/plugins/kubernetes/spot_monitor_sidecar.pyi +1 -1
  174. metaflow-stubs/plugins/ollama/__init__.pyi +1 -1
  175. metaflow-stubs/plugins/parallel_decorator.pyi +1 -1
  176. metaflow-stubs/plugins/perimeters.pyi +1 -1
  177. metaflow-stubs/plugins/project_decorator.pyi +1 -1
  178. metaflow-stubs/plugins/pypi/__init__.pyi +1 -1
  179. metaflow-stubs/plugins/pypi/conda_decorator.pyi +1 -1
  180. metaflow-stubs/plugins/pypi/conda_environment.pyi +4 -4
  181. metaflow-stubs/plugins/pypi/parsers.pyi +1 -1
  182. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +1 -1
  183. metaflow-stubs/plugins/pypi/pypi_environment.pyi +1 -1
  184. metaflow-stubs/plugins/pypi/utils.pyi +1 -1
  185. metaflow-stubs/plugins/resources_decorator.pyi +1 -1
  186. metaflow-stubs/plugins/retry_decorator.pyi +1 -1
  187. metaflow-stubs/plugins/secrets/__init__.pyi +2 -2
  188. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +3 -3
  189. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +1 -1
  190. metaflow-stubs/plugins/snowflake/__init__.pyi +1 -1
  191. metaflow-stubs/plugins/storage_executor.pyi +1 -1
  192. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +2 -2
  193. metaflow-stubs/plugins/timeout_decorator.pyi +2 -2
  194. metaflow-stubs/plugins/torchtune/__init__.pyi +1 -1
  195. metaflow-stubs/plugins/uv/__init__.pyi +1 -1
  196. metaflow-stubs/plugins/uv/uv_environment.pyi +2 -2
  197. metaflow-stubs/profilers/__init__.pyi +1 -1
  198. metaflow-stubs/pylint_wrapper.pyi +1 -1
  199. metaflow-stubs/runner/__init__.pyi +1 -1
  200. metaflow-stubs/runner/deployer.pyi +4 -4
  201. metaflow-stubs/runner/deployer_impl.pyi +2 -2
  202. metaflow-stubs/runner/metaflow_runner.pyi +3 -3
  203. metaflow-stubs/runner/nbdeploy.pyi +1 -1
  204. metaflow-stubs/runner/nbrun.pyi +1 -1
  205. metaflow-stubs/runner/subprocess_manager.pyi +1 -1
  206. metaflow-stubs/runner/utils.pyi +2 -2
  207. metaflow-stubs/system/__init__.pyi +1 -1
  208. metaflow-stubs/system/system_logger.pyi +1 -1
  209. metaflow-stubs/system/system_monitor.pyi +1 -1
  210. metaflow-stubs/tagging_util.pyi +1 -1
  211. metaflow-stubs/tuple_util.pyi +1 -1
  212. metaflow-stubs/user_configs/__init__.pyi +1 -1
  213. metaflow-stubs/user_configs/config_decorators.pyi +5 -5
  214. metaflow-stubs/user_configs/config_options.pyi +2 -2
  215. metaflow-stubs/user_configs/config_parameters.pyi +6 -6
  216. {ob_metaflow_stubs-6.0.3.181.dist-info → ob_metaflow_stubs-6.0.3.182rc1.dist-info}/METADATA +1 -1
  217. ob_metaflow_stubs-6.0.3.182rc1.dist-info/RECORD +220 -0
  218. ob_metaflow_stubs-6.0.3.181.dist-info/RECORD +0 -215
  219. {ob_metaflow_stubs-6.0.3.181.dist-info → ob_metaflow_stubs-6.0.3.182rc1.dist-info}/WHEEL +0 -0
  220. {ob_metaflow_stubs-6.0.3.181.dist-info → ob_metaflow_stubs-6.0.3.182rc1.dist-info}/top_level.txt +0 -0
@@ -1,15 +1,15 @@
1
1
  ######################################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
3
  # MF version: 2.15.17.1+obcheckpoint(0.2.1);ob(v1) #
4
- # Generated on 2025-06-18T08:04:08.962280 #
4
+ # Generated on 2025-06-18T10:06:22.722981 #
5
5
  ######################################################################################################
6
6
 
7
7
  from __future__ import annotations
8
8
 
9
9
  import typing
10
10
  if typing.TYPE_CHECKING:
11
- import typing
12
11
  import datetime
12
+ import typing
13
13
  FlowSpecDerived = typing.TypeVar("FlowSpecDerived", bound="FlowSpec", contravariant=False, covariant=False)
14
14
  StepFlag = typing.NewType("StepFlag", bool)
15
15
 
@@ -35,17 +35,17 @@ 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 metaflow_git as metaflow_git
39
+ from . import tuple_util as tuple_util
38
40
  from . import cards as cards
39
41
  from . import events as events
40
- from . import tuple_util as tuple_util
41
- from . import metaflow_git as metaflow_git
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 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 .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
50
50
  from . import client as client
51
51
  from .client.core import namespace as namespace
@@ -78,6 +78,7 @@ from . import system as system
78
78
  from . import pylint_wrapper as pylint_wrapper
79
79
  from . import cli as cli
80
80
  from . import profilers as profilers
81
+ from . import ob_internal as ob_internal
81
82
 
82
83
  EXT_PKG: str
83
84
 
@@ -205,154 +206,249 @@ def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
205
206
  """
206
207
  ...
207
208
 
208
- @typing.overload
209
- 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]]]:
209
+ 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]]]:
210
210
  """
211
- Specifies the number of times the task corresponding
212
- to a step needs to be retried.
211
+ Specifies that this step should execute on Kubernetes.
213
212
 
214
- This decorator is useful for handling transient errors, such as networking issues.
215
- If your task contains operations that can't be retried safely, e.g. database updates,
216
- it is advisable to annotate it with `@retry(times=0)`.
217
213
 
218
- This can be used in conjunction with the `@catch` decorator. The `@catch`
219
- decorator will execute a no-op task after all retries have been exhausted,
220
- ensuring that the flow execution can continue.
214
+ Parameters
215
+ ----------
216
+ cpu : int, default 1
217
+ Number of CPUs required for this step. If `@resources` is
218
+ also present, the maximum value from all decorators is used.
219
+ memory : int, default 4096
220
+ Memory size (in MB) required for this step. If
221
+ `@resources` is also present, the maximum value from all decorators is
222
+ used.
223
+ disk : int, default 10240
224
+ Disk size (in MB) required for this step. If
225
+ `@resources` is also present, the maximum value from all decorators is
226
+ used.
227
+ image : str, optional, default None
228
+ Docker image to use when launching on Kubernetes. If not specified, and
229
+ METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
230
+ not, a default Docker image mapping to the current version of Python is used.
231
+ image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
232
+ If given, the imagePullPolicy to be applied to the Docker image of the step.
233
+ image_pull_secrets: List[str], default []
234
+ The default is extracted from METAFLOW_KUBERNETES_IMAGE_PULL_SECRETS.
235
+ Kubernetes image pull secrets to use when pulling container images
236
+ in Kubernetes.
237
+ service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
238
+ Kubernetes service account to use when launching pod in Kubernetes.
239
+ secrets : List[str], optional, default None
240
+ Kubernetes secrets to use when launching pod in Kubernetes. These
241
+ secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
242
+ in Metaflow configuration.
243
+ node_selector: Union[Dict[str,str], str], optional, default None
244
+ Kubernetes node selector(s) to apply to the pod running the task.
245
+ Can be passed in as a comma separated string of values e.g.
246
+ 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
247
+ {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
248
+ namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
249
+ Kubernetes namespace to use when launching pod in Kubernetes.
250
+ gpu : int, optional, default None
251
+ Number of GPUs required for this step. A value of zero implies that
252
+ the scheduled node should not have GPUs.
253
+ gpu_vendor : str, default KUBERNETES_GPU_VENDOR
254
+ The vendor of the GPUs to be used for this step.
255
+ tolerations : List[str], default []
256
+ The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
257
+ Kubernetes tolerations to use when launching pod in Kubernetes.
258
+ labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
259
+ Kubernetes labels to use when launching pod in Kubernetes.
260
+ annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
261
+ Kubernetes annotations to use when launching pod in Kubernetes.
262
+ use_tmpfs : bool, default False
263
+ This enables an explicit tmpfs mount for this step.
264
+ tmpfs_tempdir : bool, default True
265
+ sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
266
+ tmpfs_size : int, optional, default: None
267
+ The value for the size (in MiB) of the tmpfs mount for this step.
268
+ This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
269
+ memory allocated for this step.
270
+ tmpfs_path : str, optional, default /metaflow_temp
271
+ Path to tmpfs mount for this step.
272
+ persistent_volume_claims : Dict[str, str], optional, default None
273
+ A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
274
+ volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
275
+ shared_memory: int, optional
276
+ Shared memory size (in MiB) required for this step
277
+ port: int, optional
278
+ Port number to specify in the Kubernetes job object
279
+ compute_pool : str, optional, default None
280
+ Compute pool to be used for for this step.
281
+ If not specified, any accessible compute pool within the perimeter is used.
282
+ hostname_resolution_timeout: int, default 10 * 60
283
+ Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
284
+ Only applicable when @parallel is used.
285
+ qos: str, default: Burstable
286
+ Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
287
+
288
+ security_context: Dict[str, Any], optional, default None
289
+ Container security context. Applies to the task container. Allows the following keys:
290
+ - privileged: bool, optional, default None
291
+ - allow_privilege_escalation: bool, optional, default None
292
+ - run_as_user: int, optional, default None
293
+ - run_as_group: int, optional, default None
294
+ - run_as_non_root: bool, optional, default None
295
+ """
296
+ ...
297
+
298
+ @typing.overload
299
+ 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]]]:
300
+ """
301
+ Creates a human-readable report, a Metaflow Card, after this step completes.
302
+
303
+ Note that you may add multiple `@card` decorators in a step with different parameters.
221
304
 
222
305
 
223
306
  Parameters
224
307
  ----------
225
- times : int, default 3
226
- Number of times to retry this task.
227
- minutes_between_retries : int, default 2
228
- Number of minutes between retries.
308
+ type : str, default 'default'
309
+ Card type.
310
+ id : str, optional, default None
311
+ If multiple cards are present, use this id to identify this card.
312
+ options : Dict[str, Any], default {}
313
+ Options passed to the card. The contents depend on the card type.
314
+ timeout : int, default 45
315
+ Interrupt reporting if it takes more than this many seconds.
229
316
  """
230
317
  ...
231
318
 
232
319
  @typing.overload
233
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
320
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
234
321
  ...
235
322
 
236
323
  @typing.overload
237
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
324
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
238
325
  ...
239
326
 
240
- 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):
327
+ 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):
241
328
  """
242
- Specifies the number of times the task corresponding
243
- to a step needs to be retried.
244
-
245
- This decorator is useful for handling transient errors, such as networking issues.
246
- If your task contains operations that can't be retried safely, e.g. database updates,
247
- it is advisable to annotate it with `@retry(times=0)`.
329
+ Creates a human-readable report, a Metaflow Card, after this step completes.
248
330
 
249
- This can be used in conjunction with the `@catch` decorator. The `@catch`
250
- decorator will execute a no-op task after all retries have been exhausted,
251
- ensuring that the flow execution can continue.
331
+ Note that you may add multiple `@card` decorators in a step with different parameters.
252
332
 
253
333
 
254
334
  Parameters
255
335
  ----------
256
- times : int, default 3
257
- Number of times to retry this task.
258
- minutes_between_retries : int, default 2
259
- Number of minutes between retries.
336
+ type : str, default 'default'
337
+ Card type.
338
+ id : str, optional, default None
339
+ If multiple cards are present, use this id to identify this card.
340
+ options : Dict[str, Any], default {}
341
+ Options passed to the card. The contents depend on the card type.
342
+ timeout : int, default 45
343
+ Interrupt reporting if it takes more than this many seconds.
260
344
  """
261
345
  ...
262
346
 
263
347
  @typing.overload
264
- 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]]]:
348
+ 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]]]:
265
349
  """
266
- Specifies environment variables to be set prior to the execution of a step.
350
+ Specifies the PyPI packages for the step.
351
+
352
+ Information in this decorator will augment any
353
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
354
+ you can use `@pypi_base` to set packages required by all
355
+ steps and use `@pypi` to specify step-specific overrides.
267
356
 
268
357
 
269
358
  Parameters
270
359
  ----------
271
- vars : Dict[str, str], default {}
272
- Dictionary of environment variables to set.
360
+ packages : Dict[str, str], default: {}
361
+ Packages to use for this step. The key is the name of the package
362
+ and the value is the version to use.
363
+ python : str, optional, default: None
364
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
365
+ that the version used will correspond to the version of the Python interpreter used to start the run.
273
366
  """
274
367
  ...
275
368
 
276
369
  @typing.overload
277
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
370
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
278
371
  ...
279
372
 
280
373
  @typing.overload
281
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
374
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
282
375
  ...
283
376
 
284
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
377
+ 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):
285
378
  """
286
- Specifies environment variables to be set prior to the execution of a step.
379
+ Specifies the PyPI packages for the step.
380
+
381
+ Information in this decorator will augment any
382
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
383
+ you can use `@pypi_base` to set packages required by all
384
+ steps and use `@pypi` to specify step-specific overrides.
287
385
 
288
386
 
289
387
  Parameters
290
388
  ----------
291
- vars : Dict[str, str], default {}
292
- Dictionary of environment variables to set.
389
+ packages : Dict[str, str], default: {}
390
+ Packages to use for this step. The key is the name of the package
391
+ and the value is the version to use.
392
+ python : str, optional, default: None
393
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
394
+ that the version used will correspond to the version of the Python interpreter used to start the run.
293
395
  """
294
396
  ...
295
397
 
296
- 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]]]:
398
+ @typing.overload
399
+ 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]]]:
297
400
  """
298
- This decorator is used to run vllm APIs as Metaflow task sidecars.
299
-
300
- User code call
301
- --------------
302
- @vllm(
303
- model="...",
304
- ...
305
- )
306
-
307
- Valid backend options
308
- ---------------------
309
- - 'local': Run as a separate process on the local task machine.
310
-
311
- Valid model options
312
- -------------------
313
- Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
401
+ Enables loading / saving of models within a step.
314
402
 
315
- NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
316
- If you need multiple models, you must create multiple @vllm decorators.
317
403
 
318
404
 
319
405
  Parameters
320
406
  ----------
321
- model: str
322
- HuggingFace model identifier to be served by vLLM.
323
- backend: str
324
- Determines where and how to run the vLLM process.
325
- debug: bool
326
- Whether to turn on verbose debugging logs.
327
- kwargs : Any
328
- Any other keyword arguments are passed directly to the vLLM engine.
329
- This allows for flexible configuration of vLLM server settings.
330
- For example, `tensor_parallel_size=2`.
407
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
408
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
409
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
410
+ - `current.checkpoint`
411
+ - `current.model`
412
+ - `current.huggingface_hub`
413
+
414
+ If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
415
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
416
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
417
+
418
+ temp_dir_root : str, default: None
419
+ The root directory under which `current.model.loaded` will store loaded models
331
420
  """
332
421
  ...
333
422
 
334
- 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]]]:
423
+ @typing.overload
424
+ def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
425
+ ...
426
+
427
+ @typing.overload
428
+ def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
429
+ ...
430
+
431
+ def model(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, load: typing.Union[typing.List[str], str, typing.List[typing.Tuple[str, typing.Optional[str]]]] = None, temp_dir_root: str = None):
335
432
  """
336
- Decorator that helps cache, version and store models/datasets from huggingface hub.
433
+ Enables loading / saving of models within a step.
434
+
337
435
 
338
436
 
339
437
  Parameters
340
438
  ----------
341
- temp_dir_root : str, optional
342
- The root directory that will hold the temporary directory where objects will be downloaded.
439
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
440
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
441
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
442
+ - `current.checkpoint`
443
+ - `current.model`
444
+ - `current.huggingface_hub`
343
445
 
344
- load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
345
- The list of repos (models/datasets) to load.
346
-
347
- Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
348
-
349
- - If repo (model/dataset) is not found in the datastore:
350
- - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
351
- - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
352
- - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
446
+ If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
447
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
448
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
353
449
 
354
- - If repo is found in the datastore:
355
- - Loads it directly from datastore to local path (can be temporary directory or specified path)
450
+ temp_dir_root : str, default: None
451
+ The root directory under which `current.model.loaded` will store loaded models
356
452
  """
357
453
  ...
358
454
 
@@ -413,464 +509,256 @@ def checkpoint(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None
413
509
  """
414
510
  ...
415
511
 
416
- 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]]]:
417
- """
418
- This decorator is used to run Ollama APIs as Metaflow task sidecars.
419
-
420
- User code call
421
- --------------
422
- @ollama(
423
- models=[...],
424
- ...
425
- )
426
-
427
- Valid backend options
428
- ---------------------
429
- - 'local': Run as a separate process on the local task machine.
430
- - (TODO) 'managed': Outerbounds hosts and selects compute provider.
431
- - (TODO) 'remote': Spin up separate instance to serve Ollama models.
432
-
433
- Valid model options
434
- -------------------
435
- Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
436
-
437
-
438
- Parameters
439
- ----------
440
- models: list[str]
441
- List of Ollama containers running models in sidecars.
442
- backend: str
443
- Determines where and how to run the Ollama process.
444
- force_pull: bool
445
- Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
446
- cache_update_policy: str
447
- Cache update policy: "auto", "force", or "never".
448
- force_cache_update: bool
449
- Simple override for "force" cache update policy.
450
- debug: bool
451
- Whether to turn on verbose debugging logs.
452
- circuit_breaker_config: dict
453
- Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
454
- timeout_config: dict
455
- Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
456
- """
457
- ...
458
-
459
- 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]]]:
512
+ @typing.overload
513
+ 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]]]:
460
514
  """
461
- Specifies that this step is used to deploy an instance of the app.
462
- Requires that self.app_name, self.app_port, self.entrypoint and self.deployDir is set.
515
+ Specifies secrets to be retrieved and injected as environment variables prior to
516
+ the execution of a step.
463
517
 
464
518
 
465
519
  Parameters
466
520
  ----------
467
- app_port : int
468
- Number of GPUs to use.
469
- app_name : str
470
- Name of the app to deploy.
521
+ sources : List[Union[str, Dict[str, Any]]], default: []
522
+ List of secret specs, defining how the secrets are to be retrieved
471
523
  """
472
524
  ...
473
525
 
474
526
  @typing.overload
475
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
476
- """
477
- Decorator prototype for all step decorators. This function gets specialized
478
- and imported for all decorators types by _import_plugin_decorators().
479
- """
527
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
480
528
  ...
481
529
 
482
530
  @typing.overload
483
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
531
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
484
532
  ...
485
533
 
486
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
534
+ 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]]] = []):
487
535
  """
488
- Decorator prototype for all step decorators. This function gets specialized
489
- and imported for all decorators types by _import_plugin_decorators().
536
+ Specifies secrets to be retrieved and injected as environment variables prior to
537
+ the execution of a step.
538
+
539
+
540
+ Parameters
541
+ ----------
542
+ sources : List[Union[str, Dict[str, Any]]], default: []
543
+ List of secret specs, defining how the secrets are to be retrieved
490
544
  """
491
545
  ...
492
546
 
493
- @typing.overload
494
- 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]]]:
547
+ 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]]]:
495
548
  """
496
- Enables loading / saving of models within a step.
497
-
549
+ Specifies that this step should execute on DGX cloud.
498
550
 
499
551
 
500
552
  Parameters
501
553
  ----------
502
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
503
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
504
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
505
- - `current.checkpoint`
506
- - `current.model`
507
- - `current.huggingface_hub`
508
-
509
- If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
510
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
511
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
512
-
513
- temp_dir_root : str, default: None
514
- The root directory under which `current.model.loaded` will store loaded models
554
+ gpu : int
555
+ Number of GPUs to use.
556
+ gpu_type : str
557
+ Type of Nvidia GPU to use.
558
+ queue_timeout : int
559
+ Time to keep the job in NVCF's queue.
515
560
  """
516
561
  ...
517
562
 
518
- @typing.overload
519
- def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
520
- ...
521
-
522
- @typing.overload
523
- def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
524
- ...
525
-
526
- def model(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, load: typing.Union[typing.List[str], str, typing.List[typing.Tuple[str, typing.Optional[str]]]] = None, temp_dir_root: str = None):
563
+ def huggingface_hub(*, temp_dir_root: typing.Optional[str] = None, load: typing.Union[typing.List[str], typing.List[typing.Tuple[typing.Dict, str]], typing.List[typing.Tuple[str, str]], typing.List[typing.Dict], None]) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
527
564
  """
528
- Enables loading / saving of models within a step.
529
-
565
+ Decorator that helps cache, version and store models/datasets from huggingface hub.
530
566
 
531
567
 
532
568
  Parameters
533
569
  ----------
534
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
535
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
536
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
537
- - `current.checkpoint`
538
- - `current.model`
539
- - `current.huggingface_hub`
570
+ temp_dir_root : str, optional
571
+ The root directory that will hold the temporary directory where objects will be downloaded.
540
572
 
541
- If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
542
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
543
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
573
+ load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
574
+ The list of repos (models/datasets) to load.
544
575
 
545
- temp_dir_root : str, default: None
546
- The root directory under which `current.model.loaded` will store loaded models
576
+ Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
577
+
578
+ - If repo (model/dataset) is not found in the datastore:
579
+ - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
580
+ - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
581
+ - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
582
+
583
+ - If repo is found in the datastore:
584
+ - Loads it directly from datastore to local path (can be temporary directory or specified path)
547
585
  """
548
586
  ...
549
587
 
550
588
  @typing.overload
551
- 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]]]:
589
+ 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]]]:
552
590
  """
553
- Specifies secrets to be retrieved and injected as environment variables prior to
554
- the execution of a step.
591
+ Specifies the resources needed when executing this step.
592
+
593
+ Use `@resources` to specify the resource requirements
594
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
595
+
596
+ You can choose the compute layer on the command line by executing e.g.
597
+ ```
598
+ python myflow.py run --with batch
599
+ ```
600
+ or
601
+ ```
602
+ python myflow.py run --with kubernetes
603
+ ```
604
+ which executes the flow on the desired system using the
605
+ requirements specified in `@resources`.
555
606
 
556
607
 
557
608
  Parameters
558
609
  ----------
559
- sources : List[Union[str, Dict[str, Any]]], default: []
560
- List of secret specs, defining how the secrets are to be retrieved
610
+ cpu : int, default 1
611
+ Number of CPUs required for this step.
612
+ gpu : int, optional, default None
613
+ Number of GPUs required for this step.
614
+ disk : int, optional, default None
615
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
616
+ memory : int, default 4096
617
+ Memory size (in MB) required for this step.
618
+ shared_memory : int, optional, default None
619
+ The value for the size (in MiB) of the /dev/shm volume for this step.
620
+ This parameter maps to the `--shm-size` option in Docker.
561
621
  """
562
622
  ...
563
623
 
564
624
  @typing.overload
565
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
625
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
566
626
  ...
567
627
 
568
628
  @typing.overload
569
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
629
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
570
630
  ...
571
631
 
572
- 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]]] = []):
632
+ 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):
573
633
  """
574
- Specifies secrets to be retrieved and injected as environment variables prior to
575
- the execution of a step.
634
+ Specifies the resources needed when executing this step.
635
+
636
+ Use `@resources` to specify the resource requirements
637
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
638
+
639
+ You can choose the compute layer on the command line by executing e.g.
640
+ ```
641
+ python myflow.py run --with batch
642
+ ```
643
+ or
644
+ ```
645
+ python myflow.py run --with kubernetes
646
+ ```
647
+ which executes the flow on the desired system using the
648
+ requirements specified in `@resources`.
576
649
 
577
650
 
578
651
  Parameters
579
652
  ----------
580
- sources : List[Union[str, Dict[str, Any]]], default: []
581
- List of secret specs, defining how the secrets are to be retrieved
653
+ cpu : int, default 1
654
+ Number of CPUs required for this step.
655
+ gpu : int, optional, default None
656
+ Number of GPUs required for this step.
657
+ disk : int, optional, default None
658
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
659
+ memory : int, default 4096
660
+ Memory size (in MB) required for this step.
661
+ shared_memory : int, optional, default None
662
+ The value for the size (in MiB) of the /dev/shm volume for this step.
663
+ This parameter maps to the `--shm-size` option in Docker.
582
664
  """
583
665
  ...
584
666
 
585
667
  @typing.overload
586
- 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]]]:
668
+ 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]]]:
587
669
  """
588
- Specifies the PyPI packages for the step.
670
+ Specifies the number of times the task corresponding
671
+ to a step needs to be retried.
589
672
 
590
- Information in this decorator will augment any
591
- attributes set in the `@pyi_base` flow-level decorator. Hence,
592
- you can use `@pypi_base` to set packages required by all
593
- steps and use `@pypi` to specify step-specific overrides.
673
+ This decorator is useful for handling transient errors, such as networking issues.
674
+ If your task contains operations that can't be retried safely, e.g. database updates,
675
+ it is advisable to annotate it with `@retry(times=0)`.
676
+
677
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
678
+ decorator will execute a no-op task after all retries have been exhausted,
679
+ ensuring that the flow execution can continue.
594
680
 
595
681
 
596
682
  Parameters
597
683
  ----------
598
- packages : Dict[str, str], default: {}
599
- Packages to use for this step. The key is the name of the package
600
- and the value is the version to use.
601
- python : str, optional, default: None
602
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
603
- that the version used will correspond to the version of the Python interpreter used to start the run.
684
+ times : int, default 3
685
+ Number of times to retry this task.
686
+ minutes_between_retries : int, default 2
687
+ Number of minutes between retries.
604
688
  """
605
689
  ...
606
690
 
607
691
  @typing.overload
608
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
692
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
609
693
  ...
610
694
 
611
695
  @typing.overload
612
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
696
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
613
697
  ...
614
698
 
615
- 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):
699
+ 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):
616
700
  """
617
- Specifies the PyPI packages for the step.
701
+ Specifies the number of times the task corresponding
702
+ to a step needs to be retried.
618
703
 
619
- Information in this decorator will augment any
620
- attributes set in the `@pyi_base` flow-level decorator. Hence,
621
- you can use `@pypi_base` to set packages required by all
622
- steps and use `@pypi` to specify step-specific overrides.
623
-
624
-
625
- Parameters
626
- ----------
627
- packages : Dict[str, str], default: {}
628
- Packages to use for this step. The key is the name of the package
629
- and the value is the version to use.
630
- python : str, optional, default: None
631
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
632
- that the version used will correspond to the version of the Python interpreter used to start the run.
633
- """
634
- ...
635
-
636
- 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]]]:
637
- """
638
- Specifies that this step should execute on Kubernetes.
639
-
640
-
641
- Parameters
642
- ----------
643
- cpu : int, default 1
644
- Number of CPUs required for this step. If `@resources` is
645
- also present, the maximum value from all decorators is used.
646
- memory : int, default 4096
647
- Memory size (in MB) required for this step. If
648
- `@resources` is also present, the maximum value from all decorators is
649
- used.
650
- disk : int, default 10240
651
- Disk size (in MB) required for this step. If
652
- `@resources` is also present, the maximum value from all decorators is
653
- used.
654
- image : str, optional, default None
655
- Docker image to use when launching on Kubernetes. If not specified, and
656
- METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
657
- not, a default Docker image mapping to the current version of Python is used.
658
- image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
659
- If given, the imagePullPolicy to be applied to the Docker image of the step.
660
- image_pull_secrets: List[str], default []
661
- The default is extracted from METAFLOW_KUBERNETES_IMAGE_PULL_SECRETS.
662
- Kubernetes image pull secrets to use when pulling container images
663
- in Kubernetes.
664
- service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
665
- Kubernetes service account to use when launching pod in Kubernetes.
666
- secrets : List[str], optional, default None
667
- Kubernetes secrets to use when launching pod in Kubernetes. These
668
- secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
669
- in Metaflow configuration.
670
- node_selector: Union[Dict[str,str], str], optional, default None
671
- Kubernetes node selector(s) to apply to the pod running the task.
672
- Can be passed in as a comma separated string of values e.g.
673
- 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
674
- {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
675
- namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
676
- Kubernetes namespace to use when launching pod in Kubernetes.
677
- gpu : int, optional, default None
678
- Number of GPUs required for this step. A value of zero implies that
679
- the scheduled node should not have GPUs.
680
- gpu_vendor : str, default KUBERNETES_GPU_VENDOR
681
- The vendor of the GPUs to be used for this step.
682
- tolerations : List[str], default []
683
- The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
684
- Kubernetes tolerations to use when launching pod in Kubernetes.
685
- labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
686
- Kubernetes labels to use when launching pod in Kubernetes.
687
- annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
688
- Kubernetes annotations to use when launching pod in Kubernetes.
689
- use_tmpfs : bool, default False
690
- This enables an explicit tmpfs mount for this step.
691
- tmpfs_tempdir : bool, default True
692
- sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
693
- tmpfs_size : int, optional, default: None
694
- The value for the size (in MiB) of the tmpfs mount for this step.
695
- This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
696
- memory allocated for this step.
697
- tmpfs_path : str, optional, default /metaflow_temp
698
- Path to tmpfs mount for this step.
699
- persistent_volume_claims : Dict[str, str], optional, default None
700
- A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
701
- volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
702
- shared_memory: int, optional
703
- Shared memory size (in MiB) required for this step
704
- port: int, optional
705
- Port number to specify in the Kubernetes job object
706
- compute_pool : str, optional, default None
707
- Compute pool to be used for for this step.
708
- If not specified, any accessible compute pool within the perimeter is used.
709
- hostname_resolution_timeout: int, default 10 * 60
710
- Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
711
- Only applicable when @parallel is used.
712
- qos: str, default: Burstable
713
- Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
714
-
715
- security_context: Dict[str, Any], optional, default None
716
- Container security context. Applies to the task container. Allows the following keys:
717
- - privileged: bool, optional, default None
718
- - allow_privilege_escalation: bool, optional, default None
719
- - run_as_user: int, optional, default None
720
- - run_as_group: int, optional, default None
721
- - run_as_non_root: bool, optional, default None
722
- """
723
- ...
724
-
725
- 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]]]:
726
- """
727
- Specifies that this step should execute on DGX cloud.
728
-
729
-
730
- Parameters
731
- ----------
732
- gpu : int
733
- Number of GPUs to use.
734
- gpu_type : str
735
- Type of Nvidia GPU to use.
736
- """
737
- ...
738
-
739
- @typing.overload
740
- 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]]]:
741
- """
742
- Specifies the Conda environment for the step.
743
-
744
- Information in this decorator will augment any
745
- attributes set in the `@conda_base` flow-level decorator. Hence,
746
- you can use `@conda_base` to set packages required by all
747
- steps and use `@conda` to specify step-specific overrides.
748
-
749
-
750
- Parameters
751
- ----------
752
- packages : Dict[str, str], default {}
753
- Packages to use for this step. The key is the name of the package
754
- and the value is the version to use.
755
- libraries : Dict[str, str], default {}
756
- Supported for backward compatibility. When used with packages, packages will take precedence.
757
- python : str, optional, default None
758
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
759
- that the version used will correspond to the version of the Python interpreter used to start the run.
760
- disabled : bool, default False
761
- If set to True, disables @conda.
762
- """
763
- ...
764
-
765
- @typing.overload
766
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
767
- ...
768
-
769
- @typing.overload
770
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
771
- ...
772
-
773
- 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):
774
- """
775
- Specifies the Conda environment for the step.
704
+ This decorator is useful for handling transient errors, such as networking issues.
705
+ If your task contains operations that can't be retried safely, e.g. database updates,
706
+ it is advisable to annotate it with `@retry(times=0)`.
776
707
 
777
- Information in this decorator will augment any
778
- attributes set in the `@conda_base` flow-level decorator. Hence,
779
- you can use `@conda_base` to set packages required by all
780
- steps and use `@conda` to specify step-specific overrides.
708
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
709
+ decorator will execute a no-op task after all retries have been exhausted,
710
+ ensuring that the flow execution can continue.
781
711
 
782
712
 
783
713
  Parameters
784
714
  ----------
785
- packages : Dict[str, str], default {}
786
- Packages to use for this step. The key is the name of the package
787
- and the value is the version to use.
788
- libraries : Dict[str, str], default {}
789
- Supported for backward compatibility. When used with packages, packages will take precedence.
790
- python : str, optional, default None
791
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
792
- that the version used will correspond to the version of the Python interpreter used to start the run.
793
- disabled : bool, default False
794
- If set to True, disables @conda.
715
+ times : int, default 3
716
+ Number of times to retry this task.
717
+ minutes_between_retries : int, default 2
718
+ Number of minutes between retries.
795
719
  """
796
720
  ...
797
721
 
798
- @typing.overload
799
- 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]]]:
722
+ 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]]]:
800
723
  """
801
- Specifies the resources needed when executing this step.
802
-
803
- Use `@resources` to specify the resource requirements
804
- independently of the specific compute layer (`@batch`, `@kubernetes`).
805
-
806
- You can choose the compute layer on the command line by executing e.g.
807
- ```
808
- python myflow.py run --with batch
809
- ```
810
- or
811
- ```
812
- python myflow.py run --with kubernetes
813
- ```
814
- which executes the flow on the desired system using the
815
- requirements specified in `@resources`.
816
-
724
+ This decorator is used to run Ollama APIs as Metaflow task sidecars.
817
725
 
818
- Parameters
819
- ----------
820
- cpu : int, default 1
821
- Number of CPUs required for this step.
822
- gpu : int, optional, default None
823
- Number of GPUs required for this step.
824
- disk : int, optional, default None
825
- Disk size (in MB) required for this step. Only applies on Kubernetes.
826
- memory : int, default 4096
827
- Memory size (in MB) required for this step.
828
- shared_memory : int, optional, default None
829
- The value for the size (in MiB) of the /dev/shm volume for this step.
830
- This parameter maps to the `--shm-size` option in Docker.
831
- """
832
- ...
833
-
834
- @typing.overload
835
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
836
- ...
837
-
838
- @typing.overload
839
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
840
- ...
841
-
842
- 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):
843
- """
844
- Specifies the resources needed when executing this step.
726
+ User code call
727
+ --------------
728
+ @ollama(
729
+ models=[...],
730
+ ...
731
+ )
845
732
 
846
- Use `@resources` to specify the resource requirements
847
- independently of the specific compute layer (`@batch`, `@kubernetes`).
733
+ Valid backend options
734
+ ---------------------
735
+ - 'local': Run as a separate process on the local task machine.
736
+ - (TODO) 'managed': Outerbounds hosts and selects compute provider.
737
+ - (TODO) 'remote': Spin up separate instance to serve Ollama models.
848
738
 
849
- You can choose the compute layer on the command line by executing e.g.
850
- ```
851
- python myflow.py run --with batch
852
- ```
853
- or
854
- ```
855
- python myflow.py run --with kubernetes
856
- ```
857
- which executes the flow on the desired system using the
858
- requirements specified in `@resources`.
739
+ Valid model options
740
+ -------------------
741
+ Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
859
742
 
860
743
 
861
744
  Parameters
862
745
  ----------
863
- cpu : int, default 1
864
- Number of CPUs required for this step.
865
- gpu : int, optional, default None
866
- Number of GPUs required for this step.
867
- disk : int, optional, default None
868
- Disk size (in MB) required for this step. Only applies on Kubernetes.
869
- memory : int, default 4096
870
- Memory size (in MB) required for this step.
871
- shared_memory : int, optional, default None
872
- The value for the size (in MiB) of the /dev/shm volume for this step.
873
- This parameter maps to the `--shm-size` option in Docker.
746
+ models: list[str]
747
+ List of Ollama containers running models in sidecars.
748
+ backend: str
749
+ Determines where and how to run the Ollama process.
750
+ force_pull: bool
751
+ Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
752
+ cache_update_policy: str
753
+ Cache update policy: "auto", "force", or "never".
754
+ force_cache_update: bool
755
+ Simple override for "force" cache update policy.
756
+ debug: bool
757
+ Whether to turn on verbose debugging logs.
758
+ circuit_breaker_config: dict
759
+ Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
760
+ timeout_config: dict
761
+ Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
874
762
  """
875
763
  ...
876
764
 
@@ -934,357 +822,197 @@ def timeout(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None],
934
822
  ...
935
823
 
936
824
  @typing.overload
937
- def fast_bakery_internal(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]]]:
938
826
  """
939
- Internal decorator to support Fast bakery
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.
940
847
  """
941
848
  ...
942
849
 
943
850
  @typing.overload
944
- def fast_bakery_internal(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]:
945
852
  ...
946
853
 
947
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
948
- """
949
- Internal decorator to support Fast bakery
950
- """
854
+ @typing.overload
855
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
951
856
  ...
952
857
 
953
- 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]]]:
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):
954
859
  """
955
- Specifies that this step should execute on DGX cloud.
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.
956
866
 
957
867
 
958
868
  Parameters
959
869
  ----------
960
- gpu : int
961
- Number of GPUs to use.
962
- gpu_type : str
963
- Type of Nvidia GPU to use.
964
- queue_timeout : int
965
- Time to keep the job in NVCF's queue.
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.
966
880
  """
967
881
  ...
968
882
 
969
883
  @typing.overload
970
- 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]]]:
884
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
971
885
  """
972
- Creates a human-readable report, a Metaflow Card, after this step completes.
973
-
974
- Note that you may add multiple `@card` decorators in a step with different parameters.
975
-
976
-
977
- Parameters
978
- ----------
979
- type : str, default 'default'
980
- Card type.
981
- id : str, optional, default None
982
- If multiple cards are present, use this id to identify this card.
983
- options : Dict[str, Any], default {}
984
- Options passed to the card. The contents depend on the card type.
985
- timeout : int, default 45
986
- Interrupt reporting if it takes more than this many seconds.
886
+ Decorator prototype for all step decorators. This function gets specialized
887
+ and imported for all decorators types by _import_plugin_decorators().
987
888
  """
988
889
  ...
989
890
 
990
891
  @typing.overload
991
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
992
- ...
993
-
994
- @typing.overload
995
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
892
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
996
893
  ...
997
894
 
998
- 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):
895
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
999
896
  """
1000
- Creates a human-readable report, a Metaflow Card, after this step completes.
1001
-
1002
- Note that you may add multiple `@card` decorators in a step with different parameters.
1003
-
1004
-
1005
- Parameters
1006
- ----------
1007
- type : str, default 'default'
1008
- Card type.
1009
- id : str, optional, default None
1010
- If multiple cards are present, use this id to identify this card.
1011
- options : Dict[str, Any], default {}
1012
- Options passed to the card. The contents depend on the card type.
1013
- timeout : int, default 45
1014
- Interrupt reporting if it takes more than this many seconds.
897
+ Decorator prototype for all step decorators. This function gets specialized
898
+ and imported for all decorators types by _import_plugin_decorators().
1015
899
  """
1016
900
  ...
1017
901
 
1018
- def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
902
+ @typing.overload
903
+ 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]]]:
1019
904
  """
1020
- Specifies what flows belong to the same project.
1021
-
1022
- A project-specific namespace is created for all flows that
1023
- use the same `@project(name)`.
905
+ Specifies environment variables to be set prior to the execution of a step.
1024
906
 
1025
907
 
1026
908
  Parameters
1027
909
  ----------
1028
- name : str
1029
- Project name. Make sure that the name is unique amongst all
1030
- projects that use the same production scheduler. The name may
1031
- contain only lowercase alphanumeric characters and underscores.
1032
-
1033
- branch : Optional[str], default None
1034
- The branch to use. If not specified, the branch is set to
1035
- `user.<username>` unless `production` is set to `True`. This can
1036
- also be set on the command line using `--branch` as a top-level option.
1037
- It is an error to specify `branch` in the decorator and on the command line.
1038
-
1039
- production : bool, default False
1040
- Whether or not the branch is the production branch. This can also be set on the
1041
- command line using `--production` as a top-level option. It is an error to specify
1042
- `production` in the decorator and on the command line.
1043
- The project branch name will be:
1044
- - if `branch` is specified:
1045
- - if `production` is True: `prod.<branch>`
1046
- - if `production` is False: `test.<branch>`
1047
- - if `branch` is not specified:
1048
- - if `production` is True: `prod`
1049
- - if `production` is False: `user.<username>`
910
+ vars : Dict[str, str], default {}
911
+ Dictionary of environment variables to set.
1050
912
  """
1051
913
  ...
1052
914
 
1053
- 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]]:
1054
- """
1055
- 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.
1056
- 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.
1057
-
1058
-
1059
- Parameters
1060
- ----------
1061
- timeout : int
1062
- Time, in seconds before the task times out and fails. (Default: 3600)
1063
- poke_interval : int
1064
- Time in seconds that the job should wait in between each try. (Default: 60)
1065
- mode : str
1066
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1067
- exponential_backoff : bool
1068
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1069
- pool : str
1070
- the slot pool this task should run in,
1071
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1072
- soft_fail : bool
1073
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1074
- name : str
1075
- Name of the sensor on Airflow
1076
- description : str
1077
- Description of sensor in the Airflow UI
1078
- external_dag_id : str
1079
- The dag_id that contains the task you want to wait for.
1080
- external_task_ids : List[str]
1081
- The list of task_ids that you want to wait for.
1082
- If None (default value) the sensor waits for the DAG. (Default: None)
1083
- allowed_states : List[str]
1084
- Iterable of allowed states, (Default: ['success'])
1085
- failed_states : List[str]
1086
- Iterable of failed or dis-allowed states. (Default: None)
1087
- execution_delta : datetime.timedelta
1088
- time difference with the previous execution to look at,
1089
- the default is the same logical date as the current task or DAG. (Default: None)
1090
- check_existence: bool
1091
- Set to True to check if the external task exists or check if
1092
- the DAG to wait for exists. (Default: True)
1093
- """
915
+ @typing.overload
916
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1094
917
  ...
1095
918
 
1096
919
  @typing.overload
1097
- 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]]:
920
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
921
+ ...
922
+
923
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
1098
924
  """
1099
- Specifies the Conda environment for all steps of the flow.
1100
-
1101
- Use `@conda_base` to set common libraries required by all
1102
- steps and use `@conda` to specify step-specific additions.
925
+ Specifies environment variables to be set prior to the execution of a step.
1103
926
 
1104
927
 
1105
928
  Parameters
1106
929
  ----------
1107
- packages : Dict[str, str], default {}
1108
- Packages to use for this flow. The key is the name of the package
1109
- and the value is the version to use.
1110
- libraries : Dict[str, str], default {}
1111
- Supported for backward compatibility. When used with packages, packages will take precedence.
1112
- python : str, optional, default None
1113
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1114
- that the version used will correspond to the version of the Python interpreter used to start the run.
1115
- disabled : bool, default False
1116
- If set to True, disables Conda.
930
+ vars : Dict[str, str], default {}
931
+ Dictionary of environment variables to set.
1117
932
  """
1118
933
  ...
1119
934
 
1120
- @typing.overload
1121
- def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1122
- ...
1123
-
1124
- 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):
935
+ 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]]]:
1125
936
  """
1126
- Specifies the Conda environment for all steps of the flow.
1127
-
1128
- Use `@conda_base` to set common libraries required by all
1129
- steps and use `@conda` to specify step-specific additions.
937
+ Specifies that this step is used to deploy an instance of the app.
938
+ Requires that self.app_name, self.app_port, self.entrypoint and self.deployDir is set.
1130
939
 
1131
940
 
1132
941
  Parameters
1133
942
  ----------
1134
- packages : Dict[str, str], default {}
1135
- Packages to use for this flow. The key is the name of the package
1136
- and the value is the version to use.
1137
- libraries : Dict[str, str], default {}
1138
- Supported for backward compatibility. When used with packages, packages will take precedence.
1139
- python : str, optional, default None
1140
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1141
- that the version used will correspond to the version of the Python interpreter used to start the run.
1142
- disabled : bool, default False
1143
- If set to True, disables Conda.
943
+ app_port : int
944
+ Number of GPUs to use.
945
+ app_name : str
946
+ Name of the app to deploy.
1144
947
  """
1145
948
  ...
1146
949
 
1147
- @typing.overload
1148
- 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]]:
950
+ 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]]]:
1149
951
  """
1150
- Specifies the event(s) that this flow depends on.
1151
-
1152
- ```
1153
- @trigger(event='foo')
1154
- ```
1155
- or
1156
- ```
1157
- @trigger(events=['foo', 'bar'])
1158
- ```
1159
-
1160
- Additionally, you can specify the parameter mappings
1161
- to map event payload to Metaflow parameters for the flow.
1162
- ```
1163
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1164
- ```
1165
- or
1166
- ```
1167
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1168
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1169
- ```
1170
-
1171
- 'parameters' can also be a list of strings and tuples like so:
1172
- ```
1173
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1174
- ```
1175
- This is equivalent to:
1176
- ```
1177
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1178
- ```
952
+ Specifies that this step should execute on DGX cloud.
1179
953
 
1180
954
 
1181
955
  Parameters
1182
956
  ----------
1183
- event : Union[str, Dict[str, Any]], optional, default None
1184
- Event dependency for this flow.
1185
- events : List[Union[str, Dict[str, Any]]], default []
1186
- Events dependency for this flow.
1187
- options : Dict[str, Any], default {}
1188
- Backend-specific configuration for tuning eventing behavior.
957
+ gpu : int
958
+ Number of GPUs to use.
959
+ gpu_type : str
960
+ Type of Nvidia GPU to use.
1189
961
  """
1190
962
  ...
1191
963
 
1192
- @typing.overload
1193
- def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1194
- ...
1195
-
1196
- def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: typing.Union[str, typing.Dict[str, typing.Any], None] = None, events: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], options: typing.Dict[str, typing.Any] = {}):
964
+ 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]]]:
1197
965
  """
1198
- Specifies the event(s) that this flow depends on.
966
+ This decorator is used to run vllm APIs as Metaflow task sidecars.
1199
967
 
1200
- ```
1201
- @trigger(event='foo')
1202
- ```
1203
- or
1204
- ```
1205
- @trigger(events=['foo', 'bar'])
1206
- ```
968
+ User code call
969
+ --------------
970
+ @vllm(
971
+ model="...",
972
+ ...
973
+ )
1207
974
 
1208
- Additionally, you can specify the parameter mappings
1209
- to map event payload to Metaflow parameters for the flow.
1210
- ```
1211
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1212
- ```
1213
- or
1214
- ```
1215
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1216
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1217
- ```
975
+ Valid backend options
976
+ ---------------------
977
+ - 'local': Run as a separate process on the local task machine.
1218
978
 
1219
- 'parameters' can also be a list of strings and tuples like so:
1220
- ```
1221
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1222
- ```
1223
- This is equivalent to:
1224
- ```
1225
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1226
- ```
979
+ Valid model options
980
+ -------------------
981
+ Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
982
+
983
+ NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
984
+ If you need multiple models, you must create multiple @vllm decorators.
1227
985
 
1228
986
 
1229
987
  Parameters
1230
988
  ----------
1231
- event : Union[str, Dict[str, Any]], optional, default None
1232
- Event dependency for this flow.
1233
- events : List[Union[str, Dict[str, Any]]], default []
1234
- Events dependency for this flow.
1235
- options : Dict[str, Any], default {}
1236
- Backend-specific configuration for tuning eventing behavior.
989
+ model: str
990
+ HuggingFace model identifier to be served by vLLM.
991
+ backend: str
992
+ Determines where and how to run the vLLM process.
993
+ debug: bool
994
+ Whether to turn on verbose debugging logs.
995
+ kwargs : Any
996
+ Any other keyword arguments are passed directly to the vLLM engine.
997
+ This allows for flexible configuration of vLLM server settings.
998
+ For example, `tensor_parallel_size=2`.
1237
999
  """
1238
1000
  ...
1239
1001
 
1240
1002
  @typing.overload
1241
- 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]]:
1003
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1242
1004
  """
1243
- Specifies the times when the flow should be run when running on a
1244
- production scheduler.
1245
-
1246
-
1247
- Parameters
1248
- ----------
1249
- hourly : bool, default False
1250
- Run the workflow hourly.
1251
- daily : bool, default True
1252
- Run the workflow daily.
1253
- weekly : bool, default False
1254
- Run the workflow weekly.
1255
- cron : str, optional, default None
1256
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1257
- specified by this expression.
1258
- timezone : str, optional, default None
1259
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1260
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1005
+ Internal decorator to support Fast bakery
1261
1006
  """
1262
1007
  ...
1263
1008
 
1264
1009
  @typing.overload
1265
- def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1010
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1266
1011
  ...
1267
1012
 
1268
- 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):
1013
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1269
1014
  """
1270
- Specifies the times when the flow should be run when running on a
1271
- production scheduler.
1272
-
1273
-
1274
- Parameters
1275
- ----------
1276
- hourly : bool, default False
1277
- Run the workflow hourly.
1278
- daily : bool, default True
1279
- Run the workflow daily.
1280
- weekly : bool, default False
1281
- Run the workflow weekly.
1282
- cron : str, optional, default None
1283
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1284
- specified by this expression.
1285
- timezone : str, optional, default None
1286
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1287
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1015
+ Internal decorator to support Fast bakery
1288
1016
  """
1289
1017
  ...
1290
1018
 
@@ -1402,13 +1130,86 @@ def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None)
1402
1130
  """
1403
1131
  ...
1404
1132
 
1405
- 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]]:
1133
+ @typing.overload
1134
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1406
1135
  """
1407
- 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)
1408
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1409
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1410
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
1411
- starts only after all sensors finish.
1136
+ Specifies the PyPI packages for all steps of the flow.
1137
+
1138
+ Use `@pypi_base` to set common packages required by all
1139
+ steps and use `@pypi` to specify step-specific overrides.
1140
+
1141
+ Parameters
1142
+ ----------
1143
+ packages : Dict[str, str], default: {}
1144
+ Packages to use for this flow. The key is the name of the package
1145
+ and the value is the version to use.
1146
+ python : str, optional, default: None
1147
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1148
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1149
+ """
1150
+ ...
1151
+
1152
+ @typing.overload
1153
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1154
+ ...
1155
+
1156
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1157
+ """
1158
+ Specifies the PyPI packages for all steps of the flow.
1159
+
1160
+ Use `@pypi_base` to set common packages required by all
1161
+ steps and use `@pypi` to specify step-specific overrides.
1162
+
1163
+ Parameters
1164
+ ----------
1165
+ packages : Dict[str, str], default: {}
1166
+ Packages to use for this flow. The key is the name of the package
1167
+ and the value is the version to use.
1168
+ python : str, optional, default: None
1169
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1170
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1171
+ """
1172
+ ...
1173
+
1174
+ def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1175
+ """
1176
+ Specifies what flows belong to the same project.
1177
+
1178
+ A project-specific namespace is created for all flows that
1179
+ use the same `@project(name)`.
1180
+
1181
+
1182
+ Parameters
1183
+ ----------
1184
+ name : str
1185
+ Project name. Make sure that the name is unique amongst all
1186
+ projects that use the same production scheduler. The name may
1187
+ contain only lowercase alphanumeric characters and underscores.
1188
+
1189
+ branch : Optional[str], default None
1190
+ The branch to use. If not specified, the branch is set to
1191
+ `user.<username>` unless `production` is set to `True`. This can
1192
+ also be set on the command line using `--branch` as a top-level option.
1193
+ It is an error to specify `branch` in the decorator and on the command line.
1194
+
1195
+ production : bool, default False
1196
+ Whether or not the branch is the production branch. This can also be set on the
1197
+ command line using `--production` as a top-level option. It is an error to specify
1198
+ `production` in the decorator and on the command line.
1199
+ The project branch name will be:
1200
+ - if `branch` is specified:
1201
+ - if `production` is True: `prod.<branch>`
1202
+ - if `production` is False: `test.<branch>`
1203
+ - if `branch` is not specified:
1204
+ - if `production` is True: `prod`
1205
+ - if `production` is False: `user.<username>`
1206
+ """
1207
+ ...
1208
+
1209
+ 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]]:
1210
+ """
1211
+ 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.
1212
+ 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.
1412
1213
 
1413
1214
 
1414
1215
  Parameters
@@ -1430,59 +1231,165 @@ def airflow_s3_key_sensor(*, timeout: int, poke_interval: int, mode: str, expone
1430
1231
  Name of the sensor on Airflow
1431
1232
  description : str
1432
1233
  Description of sensor in the Airflow UI
1433
- bucket_key : Union[str, List[str]]
1434
- The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1435
- When it's specified as a full s3:// url, please leave `bucket_name` as None
1436
- bucket_name : str
1437
- Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1438
- When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1439
- wildcard_match : bool
1440
- whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1441
- aws_conn_id : str
1442
- a reference to the s3 connection on Airflow. (Default: None)
1443
- verify : bool
1444
- Whether or not to verify SSL certificates for S3 connection. (Default: None)
1234
+ external_dag_id : str
1235
+ The dag_id that contains the task you want to wait for.
1236
+ external_task_ids : List[str]
1237
+ The list of task_ids that you want to wait for.
1238
+ If None (default value) the sensor waits for the DAG. (Default: None)
1239
+ allowed_states : List[str]
1240
+ Iterable of allowed states, (Default: ['success'])
1241
+ failed_states : List[str]
1242
+ Iterable of failed or dis-allowed states. (Default: None)
1243
+ execution_delta : datetime.timedelta
1244
+ time difference with the previous execution to look at,
1245
+ the default is the same logical date as the current task or DAG. (Default: None)
1246
+ check_existence: bool
1247
+ Set to True to check if the external task exists or check if
1248
+ the DAG to wait for exists. (Default: True)
1445
1249
  """
1446
1250
  ...
1447
1251
 
1448
1252
  @typing.overload
1449
- def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1253
+ 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]]:
1450
1254
  """
1451
- Specifies the PyPI packages for all steps of the flow.
1255
+ Specifies the event(s) that this flow depends on.
1256
+
1257
+ ```
1258
+ @trigger(event='foo')
1259
+ ```
1260
+ or
1261
+ ```
1262
+ @trigger(events=['foo', 'bar'])
1263
+ ```
1264
+
1265
+ Additionally, you can specify the parameter mappings
1266
+ to map event payload to Metaflow parameters for the flow.
1267
+ ```
1268
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1269
+ ```
1270
+ or
1271
+ ```
1272
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1273
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1274
+ ```
1275
+
1276
+ 'parameters' can also be a list of strings and tuples like so:
1277
+ ```
1278
+ @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1279
+ ```
1280
+ This is equivalent to:
1281
+ ```
1282
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1283
+ ```
1452
1284
 
1453
- Use `@pypi_base` to set common packages required by all
1454
- steps and use `@pypi` to specify step-specific overrides.
1455
1285
 
1456
1286
  Parameters
1457
1287
  ----------
1458
- packages : Dict[str, str], default: {}
1288
+ event : Union[str, Dict[str, Any]], optional, default None
1289
+ Event dependency for this flow.
1290
+ events : List[Union[str, Dict[str, Any]]], default []
1291
+ Events dependency for this flow.
1292
+ options : Dict[str, Any], default {}
1293
+ Backend-specific configuration for tuning eventing behavior.
1294
+ """
1295
+ ...
1296
+
1297
+ @typing.overload
1298
+ def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1299
+ ...
1300
+
1301
+ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: typing.Union[str, typing.Dict[str, typing.Any], None] = None, events: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], options: typing.Dict[str, typing.Any] = {}):
1302
+ """
1303
+ Specifies the event(s) that this flow depends on.
1304
+
1305
+ ```
1306
+ @trigger(event='foo')
1307
+ ```
1308
+ or
1309
+ ```
1310
+ @trigger(events=['foo', 'bar'])
1311
+ ```
1312
+
1313
+ Additionally, you can specify the parameter mappings
1314
+ to map event payload to Metaflow parameters for the flow.
1315
+ ```
1316
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1317
+ ```
1318
+ or
1319
+ ```
1320
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1321
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1322
+ ```
1323
+
1324
+ 'parameters' can also be a list of strings and tuples like so:
1325
+ ```
1326
+ @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1327
+ ```
1328
+ This is equivalent to:
1329
+ ```
1330
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1331
+ ```
1332
+
1333
+
1334
+ Parameters
1335
+ ----------
1336
+ event : Union[str, Dict[str, Any]], optional, default None
1337
+ Event dependency for this flow.
1338
+ events : List[Union[str, Dict[str, Any]]], default []
1339
+ Events dependency for this flow.
1340
+ options : Dict[str, Any], default {}
1341
+ Backend-specific configuration for tuning eventing behavior.
1342
+ """
1343
+ ...
1344
+
1345
+ @typing.overload
1346
+ 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]]:
1347
+ """
1348
+ Specifies the Conda environment for all steps of the flow.
1349
+
1350
+ Use `@conda_base` to set common libraries required by all
1351
+ steps and use `@conda` to specify step-specific additions.
1352
+
1353
+
1354
+ Parameters
1355
+ ----------
1356
+ packages : Dict[str, str], default {}
1459
1357
  Packages to use for this flow. The key is the name of the package
1460
1358
  and the value is the version to use.
1461
- python : str, optional, default: None
1359
+ libraries : Dict[str, str], default {}
1360
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1361
+ python : str, optional, default None
1462
1362
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
1463
1363
  that the version used will correspond to the version of the Python interpreter used to start the run.
1364
+ disabled : bool, default False
1365
+ If set to True, disables Conda.
1464
1366
  """
1465
1367
  ...
1466
1368
 
1467
1369
  @typing.overload
1468
- def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1370
+ def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1469
1371
  ...
1470
1372
 
1471
- def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1373
+ 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):
1472
1374
  """
1473
- Specifies the PyPI packages for all steps of the flow.
1375
+ Specifies the Conda environment for all steps of the flow.
1376
+
1377
+ Use `@conda_base` to set common libraries required by all
1378
+ steps and use `@conda` to specify step-specific additions.
1474
1379
 
1475
- Use `@pypi_base` to set common packages required by all
1476
- steps and use `@pypi` to specify step-specific overrides.
1477
1380
 
1478
1381
  Parameters
1479
1382
  ----------
1480
- packages : Dict[str, str], default: {}
1383
+ packages : Dict[str, str], default {}
1481
1384
  Packages to use for this flow. The key is the name of the package
1482
1385
  and the value is the version to use.
1483
- python : str, optional, default: None
1386
+ libraries : Dict[str, str], default {}
1387
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1388
+ python : str, optional, default None
1484
1389
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
1485
1390
  that the version used will correspond to the version of the Python interpreter used to start the run.
1391
+ disabled : bool, default False
1392
+ If set to True, disables Conda.
1486
1393
  """
1487
1394
  ...
1488
1395
 
@@ -1587,5 +1494,99 @@ def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *
1587
1494
  """
1588
1495
  ...
1589
1496
 
1497
+ @typing.overload
1498
+ 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]]:
1499
+ """
1500
+ Specifies the times when the flow should be run when running on a
1501
+ production scheduler.
1502
+
1503
+
1504
+ Parameters
1505
+ ----------
1506
+ hourly : bool, default False
1507
+ Run the workflow hourly.
1508
+ daily : bool, default True
1509
+ Run the workflow daily.
1510
+ weekly : bool, default False
1511
+ Run the workflow weekly.
1512
+ cron : str, optional, default None
1513
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1514
+ specified by this expression.
1515
+ timezone : str, optional, default None
1516
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1517
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1518
+ """
1519
+ ...
1520
+
1521
+ @typing.overload
1522
+ def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1523
+ ...
1524
+
1525
+ 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):
1526
+ """
1527
+ Specifies the times when the flow should be run when running on a
1528
+ production scheduler.
1529
+
1530
+
1531
+ Parameters
1532
+ ----------
1533
+ hourly : bool, default False
1534
+ Run the workflow hourly.
1535
+ daily : bool, default True
1536
+ Run the workflow daily.
1537
+ weekly : bool, default False
1538
+ Run the workflow weekly.
1539
+ cron : str, optional, default None
1540
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1541
+ specified by this expression.
1542
+ timezone : str, optional, default None
1543
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1544
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1545
+ """
1546
+ ...
1547
+
1548
+ 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]]:
1549
+ """
1550
+ 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)
1551
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1552
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1553
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
1554
+ starts only after all sensors finish.
1555
+
1556
+
1557
+ Parameters
1558
+ ----------
1559
+ timeout : int
1560
+ Time, in seconds before the task times out and fails. (Default: 3600)
1561
+ poke_interval : int
1562
+ Time in seconds that the job should wait in between each try. (Default: 60)
1563
+ mode : str
1564
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1565
+ exponential_backoff : bool
1566
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1567
+ pool : str
1568
+ the slot pool this task should run in,
1569
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1570
+ soft_fail : bool
1571
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1572
+ name : str
1573
+ Name of the sensor on Airflow
1574
+ description : str
1575
+ Description of sensor in the Airflow UI
1576
+ bucket_key : Union[str, List[str]]
1577
+ The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1578
+ When it's specified as a full s3:// url, please leave `bucket_name` as None
1579
+ bucket_name : str
1580
+ Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1581
+ When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1582
+ wildcard_match : bool
1583
+ whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1584
+ aws_conn_id : str
1585
+ a reference to the s3 connection on Airflow. (Default: None)
1586
+ verify : bool
1587
+ Whether or not to verify SSL certificates for S3 connection. (Default: None)
1588
+ """
1589
+ ...
1590
+
1590
1591
  pkg_name: str
1591
1592