ob-metaflow-stubs 6.0.3.171__py2.py3-none-any.whl → 6.0.3.172__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 (209) hide show
  1. metaflow-stubs/__init__.pyi +643 -643
  2. metaflow-stubs/cards.pyi +1 -1
  3. metaflow-stubs/cli.pyi +1 -1
  4. metaflow-stubs/cli_components/__init__.pyi +1 -1
  5. metaflow-stubs/cli_components/utils.pyi +1 -1
  6. metaflow-stubs/client/__init__.pyi +1 -1
  7. metaflow-stubs/client/core.pyi +4 -4
  8. metaflow-stubs/client/filecache.pyi +2 -2
  9. metaflow-stubs/events.pyi +2 -2
  10. metaflow-stubs/exception.pyi +1 -1
  11. metaflow-stubs/flowspec.pyi +5 -5
  12. metaflow-stubs/generated_for.txt +1 -1
  13. metaflow-stubs/includefile.pyi +3 -3
  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 +2 -2
  18. metaflow-stubs/metadata_provider/util.pyi +1 -1
  19. metaflow-stubs/metaflow_config.pyi +1 -1
  20. metaflow-stubs/metaflow_current.pyi +122 -122
  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 +1 -1
  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 +2 -2
  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 +4 -4
  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 +4 -4
  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 +4 -4
  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 +3 -3
  43. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/core.pyi +1 -1
  44. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/decorator.pyi +1 -1
  45. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/exceptions.pyi +1 -1
  46. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/task_utils.pyi +3 -3
  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 +2 -2
  49. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/exceptions.pyi +1 -1
  50. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/__init__.pyi +1 -1
  51. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/decorator.pyi +2 -2
  52. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/__init__.pyi +1 -1
  53. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/core.pyi +1 -1
  54. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/exceptions.pyi +1 -1
  55. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/model_storage.pyi +3 -3
  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 +2 -2
  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 +2 -2
  63. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/tar_utils.pyi +2 -2
  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/card_utilities/__init__.pyi +1 -1
  67. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/injector.pyi +1 -1
  68. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/__init__.pyi +1 -1
  69. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/coreweave.pyi +1 -1
  70. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/nebius.pyi +1 -1
  71. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/__init__.pyi +1 -1
  72. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/constants.pyi +1 -1
  73. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/exceptions.pyi +1 -1
  74. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/ollama.pyi +1 -1
  75. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/__init__.pyi +1 -1
  76. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/snowflake.pyi +1 -1
  77. metaflow-stubs/mf_extensions/outerbounds/profilers/__init__.pyi +1 -1
  78. metaflow-stubs/mf_extensions/outerbounds/profilers/gpu.pyi +1 -1
  79. metaflow-stubs/mf_extensions/outerbounds/remote_config.pyi +1 -1
  80. metaflow-stubs/mf_extensions/outerbounds/toplevel/__init__.pyi +1 -1
  81. metaflow-stubs/mf_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.pyi +1 -1
  82. metaflow-stubs/multicore_utils.pyi +1 -1
  83. metaflow-stubs/parameters.pyi +4 -4
  84. metaflow-stubs/plugins/__init__.pyi +11 -10
  85. metaflow-stubs/plugins/airflow/__init__.pyi +1 -1
  86. metaflow-stubs/plugins/airflow/airflow_utils.pyi +1 -1
  87. metaflow-stubs/plugins/airflow/exception.pyi +1 -1
  88. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +1 -1
  89. metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi +1 -1
  90. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +1 -1
  91. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +1 -1
  92. metaflow-stubs/plugins/argo/__init__.pyi +1 -1
  93. metaflow-stubs/plugins/argo/argo_client.pyi +1 -1
  94. metaflow-stubs/plugins/argo/argo_events.pyi +1 -1
  95. metaflow-stubs/plugins/argo/argo_workflows.pyi +2 -2
  96. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +2 -2
  97. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +3 -3
  98. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +2 -2
  99. metaflow-stubs/plugins/aws/__init__.pyi +1 -1
  100. metaflow-stubs/plugins/aws/aws_client.pyi +1 -1
  101. metaflow-stubs/plugins/aws/aws_utils.pyi +1 -1
  102. metaflow-stubs/plugins/aws/batch/__init__.pyi +1 -1
  103. metaflow-stubs/plugins/aws/batch/batch.pyi +1 -1
  104. metaflow-stubs/plugins/aws/batch/batch_client.pyi +1 -1
  105. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +1 -1
  106. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +1 -1
  107. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +3 -3
  108. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +1 -1
  109. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +1 -1
  110. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +1 -1
  111. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +1 -1
  112. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +1 -1
  113. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer.pyi +3 -3
  114. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer_objects.pyi +3 -3
  115. metaflow-stubs/plugins/azure/__init__.pyi +1 -1
  116. metaflow-stubs/plugins/azure/azure_credential.pyi +1 -1
  117. metaflow-stubs/plugins/azure/azure_exceptions.pyi +1 -1
  118. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +3 -3
  119. metaflow-stubs/plugins/azure/azure_utils.pyi +1 -1
  120. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +1 -1
  121. metaflow-stubs/plugins/azure/includefile_support.pyi +1 -1
  122. metaflow-stubs/plugins/cards/__init__.pyi +1 -1
  123. metaflow-stubs/plugins/cards/card_client.pyi +2 -2
  124. metaflow-stubs/plugins/cards/card_creator.pyi +1 -1
  125. metaflow-stubs/plugins/cards/card_datastore.pyi +1 -1
  126. metaflow-stubs/plugins/cards/card_decorator.pyi +1 -1
  127. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +1 -1
  128. metaflow-stubs/plugins/cards/card_modules/basic.pyi +2 -2
  129. metaflow-stubs/plugins/cards/card_modules/card.pyi +1 -1
  130. metaflow-stubs/plugins/cards/card_modules/components.pyi +2 -2
  131. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +1 -1
  132. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +1 -1
  133. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +1 -1
  134. metaflow-stubs/plugins/cards/card_resolver.pyi +1 -1
  135. metaflow-stubs/plugins/cards/component_serializer.pyi +1 -1
  136. metaflow-stubs/plugins/cards/exception.pyi +1 -1
  137. metaflow-stubs/plugins/catch_decorator.pyi +2 -2
  138. metaflow-stubs/plugins/datatools/__init__.pyi +1 -1
  139. metaflow-stubs/plugins/datatools/local.pyi +1 -1
  140. metaflow-stubs/plugins/datatools/s3/__init__.pyi +1 -1
  141. metaflow-stubs/plugins/datatools/s3/s3.pyi +4 -4
  142. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +1 -1
  143. metaflow-stubs/plugins/datatools/s3/s3util.pyi +1 -1
  144. metaflow-stubs/plugins/debug_logger.pyi +1 -1
  145. metaflow-stubs/plugins/debug_monitor.pyi +1 -1
  146. metaflow-stubs/plugins/environment_decorator.pyi +1 -1
  147. metaflow-stubs/plugins/events_decorator.pyi +1 -1
  148. metaflow-stubs/plugins/frameworks/__init__.pyi +1 -1
  149. metaflow-stubs/plugins/frameworks/pytorch.pyi +1 -1
  150. metaflow-stubs/plugins/gcp/__init__.pyi +1 -1
  151. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +3 -3
  152. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +1 -1
  153. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +1 -1
  154. metaflow-stubs/plugins/gcp/gs_utils.pyi +1 -1
  155. metaflow-stubs/plugins/gcp/includefile_support.pyi +1 -1
  156. metaflow-stubs/plugins/kubernetes/__init__.pyi +1 -1
  157. metaflow-stubs/plugins/kubernetes/kube_utils.pyi +2 -2
  158. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +1 -1
  159. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +1 -1
  160. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +1 -1
  161. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +1 -1
  162. metaflow-stubs/plugins/kubernetes/spot_monitor_sidecar.pyi +1 -1
  163. metaflow-stubs/plugins/ollama/__init__.pyi +2 -2
  164. metaflow-stubs/plugins/parallel_decorator.pyi +1 -1
  165. metaflow-stubs/plugins/perimeters.pyi +1 -1
  166. metaflow-stubs/plugins/project_decorator.pyi +1 -1
  167. metaflow-stubs/plugins/pypi/__init__.pyi +2 -2
  168. metaflow-stubs/plugins/pypi/conda_decorator.pyi +1 -1
  169. metaflow-stubs/plugins/pypi/conda_environment.pyi +5 -5
  170. metaflow-stubs/plugins/pypi/parsers.pyi +1 -1
  171. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +1 -1
  172. metaflow-stubs/plugins/pypi/pypi_environment.pyi +1 -1
  173. metaflow-stubs/plugins/pypi/utils.pyi +1 -1
  174. metaflow-stubs/plugins/resources_decorator.pyi +1 -1
  175. metaflow-stubs/plugins/retry_decorator.pyi +1 -1
  176. metaflow-stubs/plugins/secrets/__init__.pyi +2 -2
  177. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +3 -3
  178. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +1 -1
  179. metaflow-stubs/plugins/snowflake/__init__.pyi +1 -1
  180. metaflow-stubs/plugins/storage_executor.pyi +1 -1
  181. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +2 -2
  182. metaflow-stubs/plugins/timeout_decorator.pyi +2 -2
  183. metaflow-stubs/plugins/torchtune/__init__.pyi +39 -0
  184. metaflow-stubs/plugins/uv/__init__.pyi +1 -1
  185. metaflow-stubs/plugins/uv/uv_environment.pyi +2 -2
  186. metaflow-stubs/profilers/__init__.pyi +1 -1
  187. metaflow-stubs/pylint_wrapper.pyi +1 -1
  188. metaflow-stubs/runner/__init__.pyi +1 -1
  189. metaflow-stubs/runner/deployer.pyi +5 -5
  190. metaflow-stubs/runner/deployer_impl.pyi +1 -1
  191. metaflow-stubs/runner/metaflow_runner.pyi +2 -2
  192. metaflow-stubs/runner/nbdeploy.pyi +1 -1
  193. metaflow-stubs/runner/nbrun.pyi +1 -1
  194. metaflow-stubs/runner/subprocess_manager.pyi +1 -1
  195. metaflow-stubs/runner/utils.pyi +1 -1
  196. metaflow-stubs/system/__init__.pyi +1 -1
  197. metaflow-stubs/system/system_logger.pyi +2 -2
  198. metaflow-stubs/system/system_monitor.pyi +1 -1
  199. metaflow-stubs/tagging_util.pyi +1 -1
  200. metaflow-stubs/tuple_util.pyi +1 -1
  201. metaflow-stubs/user_configs/__init__.pyi +1 -1
  202. metaflow-stubs/user_configs/config_decorators.pyi +5 -5
  203. metaflow-stubs/user_configs/config_options.pyi +3 -3
  204. metaflow-stubs/user_configs/config_parameters.pyi +5 -5
  205. {ob_metaflow_stubs-6.0.3.171.dist-info → ob_metaflow_stubs-6.0.3.172.dist-info}/METADATA +1 -1
  206. ob_metaflow_stubs-6.0.3.172.dist-info/RECORD +209 -0
  207. ob_metaflow_stubs-6.0.3.171.dist-info/RECORD +0 -208
  208. {ob_metaflow_stubs-6.0.3.171.dist-info → ob_metaflow_stubs-6.0.3.172.dist-info}/WHEEL +0 -0
  209. {ob_metaflow_stubs-6.0.3.171.dist-info → ob_metaflow_stubs-6.0.3.172.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.14.1+obcheckpoint(0.2.1);ob(v1) #
4
- # Generated on 2025-05-27T00:40:42.234892 #
4
+ # Generated on 2025-05-27T03:43:23.657123 #
5
5
  ######################################################################################################
6
6
 
7
7
  from __future__ import annotations
8
8
 
9
9
  import typing
10
10
  if typing.TYPE_CHECKING:
11
- import datetime
12
11
  import typing
12
+ import datetime
13
13
  FlowSpecDerived = typing.TypeVar("FlowSpecDerived", bound="FlowSpec", contravariant=False, covariant=False)
14
14
  StepFlag = typing.NewType("StepFlag", bool)
15
15
 
@@ -35,18 +35,18 @@ from .user_configs.config_parameters import ConfigValue as ConfigValue
35
35
  from .user_configs.config_parameters import config_expr as config_expr
36
36
  from .user_configs.config_decorators import CustomFlowDecorator as CustomFlowDecorator
37
37
  from .user_configs.config_decorators import CustomStepDecorator as CustomStepDecorator
38
+ from . import metaflow_git as metaflow_git
38
39
  from . import tuple_util as tuple_util
39
40
  from . import cards as cards
40
41
  from . import events as events
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 requirements_txt_parser as requirements_txt_parser
48
- from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
49
47
  from .plugins.pypi.parsers import conda_environment_yml_parser as conda_environment_yml_parser
48
+ from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
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
52
52
  from .client.core import get_namespace as get_namespace
@@ -154,92 +154,61 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
154
154
  ...
155
155
 
156
156
  @typing.overload
157
- 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]]]:
157
+ 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]]]:
158
158
  """
159
- Specifies the number of times the task corresponding
160
- to a step needs to be retried.
161
-
162
- This decorator is useful for handling transient errors, such as networking issues.
163
- If your task contains operations that can't be retried safely, e.g. database updates,
164
- it is advisable to annotate it with `@retry(times=0)`.
159
+ Specifies the Conda environment for the step.
165
160
 
166
- This can be used in conjunction with the `@catch` decorator. The `@catch`
167
- decorator will execute a no-op task after all retries have been exhausted,
168
- ensuring that the flow execution can continue.
161
+ Information in this decorator will augment any
162
+ attributes set in the `@conda_base` flow-level decorator. Hence,
163
+ you can use `@conda_base` to set packages required by all
164
+ steps and use `@conda` to specify step-specific overrides.
169
165
 
170
166
 
171
167
  Parameters
172
168
  ----------
173
- times : int, default 3
174
- Number of times to retry this task.
175
- minutes_between_retries : int, default 2
176
- Number of minutes between retries.
169
+ packages : Dict[str, str], default {}
170
+ Packages to use for this step. The key is the name of the package
171
+ and the value is the version to use.
172
+ libraries : Dict[str, str], default {}
173
+ Supported for backward compatibility. When used with packages, packages will take precedence.
174
+ python : str, optional, default None
175
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
176
+ that the version used will correspond to the version of the Python interpreter used to start the run.
177
+ disabled : bool, default False
178
+ If set to True, disables @conda.
177
179
  """
178
180
  ...
179
181
 
180
182
  @typing.overload
181
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
183
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
182
184
  ...
183
185
 
184
186
  @typing.overload
185
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
186
- ...
187
-
188
- 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):
189
- """
190
- Specifies the number of times the task corresponding
191
- to a step needs to be retried.
192
-
193
- This decorator is useful for handling transient errors, such as networking issues.
194
- If your task contains operations that can't be retried safely, e.g. database updates,
195
- it is advisable to annotate it with `@retry(times=0)`.
196
-
197
- This can be used in conjunction with the `@catch` decorator. The `@catch`
198
- decorator will execute a no-op task after all retries have been exhausted,
199
- ensuring that the flow execution can continue.
200
-
201
-
202
- Parameters
203
- ----------
204
- times : int, default 3
205
- Number of times to retry this task.
206
- minutes_between_retries : int, default 2
207
- Number of minutes between retries.
208
- """
187
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
209
188
  ...
210
189
 
211
- @typing.overload
212
- 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]]]:
190
+ 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):
213
191
  """
214
- Specifies secrets to be retrieved and injected as environment variables prior to
215
- the execution of a step.
216
-
192
+ Specifies the Conda environment for the step.
217
193
 
218
- Parameters
219
- ----------
220
- sources : List[Union[str, Dict[str, Any]]], default: []
221
- List of secret specs, defining how the secrets are to be retrieved
222
- """
223
- ...
224
-
225
- @typing.overload
226
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
227
- ...
228
-
229
- @typing.overload
230
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
231
- ...
232
-
233
- 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]]] = []):
234
- """
235
- Specifies secrets to be retrieved and injected as environment variables prior to
236
- the execution of a step.
194
+ Information in this decorator will augment any
195
+ attributes set in the `@conda_base` flow-level decorator. Hence,
196
+ you can use `@conda_base` to set packages required by all
197
+ steps and use `@conda` to specify step-specific overrides.
237
198
 
238
199
 
239
200
  Parameters
240
201
  ----------
241
- sources : List[Union[str, Dict[str, Any]]], default: []
242
- List of secret specs, defining how the secrets are to be retrieved
202
+ packages : Dict[str, str], default {}
203
+ Packages to use for this step. The key is the name of the package
204
+ and the value is the version to use.
205
+ libraries : Dict[str, str], default {}
206
+ Supported for backward compatibility. When used with packages, packages will take precedence.
207
+ python : str, optional, default None
208
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
209
+ that the version used will correspond to the version of the Python interpreter used to start the run.
210
+ disabled : bool, default False
211
+ If set to True, disables @conda.
243
212
  """
244
213
  ...
245
214
 
@@ -323,222 +292,199 @@ def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None]
323
292
  ...
324
293
 
325
294
  @typing.overload
326
- def timeout(*, seconds: int = 0, minutes: int = 0, hours: int = 0) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
295
+ 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]]]:
327
296
  """
328
- Specifies a timeout for your step.
329
-
330
- This decorator is useful if this step may hang indefinitely.
331
-
332
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
333
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
334
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
297
+ Enables loading / saving of models within a step.
335
298
 
336
- Note that all the values specified in parameters are added together so if you specify
337
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
338
299
 
339
300
 
340
301
  Parameters
341
302
  ----------
342
- seconds : int, default 0
343
- Number of seconds to wait prior to timing out.
344
- minutes : int, default 0
345
- Number of minutes to wait prior to timing out.
346
- hours : int, default 0
347
- Number of hours to wait prior to timing out.
303
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
304
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
305
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
306
+ - `current.checkpoint`
307
+ - `current.model`
308
+ - `current.huggingface_hub`
309
+
310
+ 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
311
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
312
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
313
+
314
+ temp_dir_root : str, default: None
315
+ The root directory under which `current.model.loaded` will store loaded models
348
316
  """
349
317
  ...
350
318
 
351
319
  @typing.overload
352
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
320
+ def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
353
321
  ...
354
322
 
355
323
  @typing.overload
356
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
324
+ def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
357
325
  ...
358
326
 
359
- def timeout(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, seconds: int = 0, minutes: int = 0, hours: int = 0):
327
+ 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):
360
328
  """
361
- Specifies a timeout for your step.
362
-
363
- This decorator is useful if this step may hang indefinitely.
364
-
365
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
366
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
367
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
329
+ Enables loading / saving of models within a step.
368
330
 
369
- Note that all the values specified in parameters are added together so if you specify
370
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
371
331
 
372
332
 
373
333
  Parameters
374
334
  ----------
375
- seconds : int, default 0
376
- Number of seconds to wait prior to timing out.
377
- minutes : int, default 0
378
- Number of minutes to wait prior to timing out.
379
- hours : int, default 0
380
- Number of hours to wait prior to timing out.
335
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
336
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
337
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
338
+ - `current.checkpoint`
339
+ - `current.model`
340
+ - `current.huggingface_hub`
341
+
342
+ 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
343
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
344
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
345
+
346
+ temp_dir_root : str, default: None
347
+ The root directory under which `current.model.loaded` will store loaded models
381
348
  """
382
349
  ...
383
350
 
384
- 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]]]:
351
+ @typing.overload
352
+ 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]]]:
385
353
  """
386
- Decorator that helps cache, version and store models/datasets from huggingface hub.
387
-
388
-
389
- Parameters
390
- ----------
391
- temp_dir_root : str, optional
392
- The root directory that will hold the temporary directory where objects will be downloaded.
354
+ Specifies the number of times the task corresponding
355
+ to a step needs to be retried.
393
356
 
394
- load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
395
- The list of repos (models/datasets) to load.
357
+ This decorator is useful for handling transient errors, such as networking issues.
358
+ If your task contains operations that can't be retried safely, e.g. database updates,
359
+ it is advisable to annotate it with `@retry(times=0)`.
396
360
 
397
- Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
361
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
362
+ decorator will execute a no-op task after all retries have been exhausted,
363
+ ensuring that the flow execution can continue.
398
364
 
399
- - If repo (model/dataset) is not found in the datastore:
400
- - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
401
- - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
402
- - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
403
365
 
404
- - If repo is found in the datastore:
405
- - Loads it directly from datastore to local path (can be temporary directory or specified path)
366
+ Parameters
367
+ ----------
368
+ times : int, default 3
369
+ Number of times to retry this task.
370
+ minutes_between_retries : int, default 2
371
+ Number of minutes between retries.
406
372
  """
407
373
  ...
408
374
 
409
- 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]]]:
375
+ @typing.overload
376
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
377
+ ...
378
+
379
+ @typing.overload
380
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
381
+ ...
382
+
383
+ 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):
410
384
  """
411
- Specifies that this step is used to deploy an instance of the app.
412
- Requires that self.app_name, self.app_port, self.entrypoint and self.deployDir is set.
385
+ Specifies the number of times the task corresponding
386
+ to a step needs to be retried.
387
+
388
+ This decorator is useful for handling transient errors, such as networking issues.
389
+ If your task contains operations that can't be retried safely, e.g. database updates,
390
+ it is advisable to annotate it with `@retry(times=0)`.
391
+
392
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
393
+ decorator will execute a no-op task after all retries have been exhausted,
394
+ ensuring that the flow execution can continue.
413
395
 
414
396
 
415
397
  Parameters
416
398
  ----------
417
- app_port : int
418
- Number of GPUs to use.
419
- app_name : str
420
- Name of the app to deploy.
399
+ times : int, default 3
400
+ Number of times to retry this task.
401
+ minutes_between_retries : int, default 2
402
+ Number of minutes between retries.
421
403
  """
422
404
  ...
423
405
 
424
- def ollama(*, models: list, backend: str, force_pull: bool, skip_push_check: bool, debug: bool) -> 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]]]:
406
+ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: typing.Optional[str] = None, image_pull_policy: str = 'KUBERNETES_IMAGE_PULL_POLICY', service_account: str = 'METAFLOW_KUBERNETES_SERVICE_ACCOUNT', secrets: typing.Optional[typing.List[str]] = None, node_selector: typing.Union[typing.Dict[str, str], str, None] = None, namespace: str = 'METAFLOW_KUBERNETES_NAMESPACE', gpu: typing.Optional[int] = None, gpu_vendor: str = 'KUBERNETES_GPU_VENDOR', tolerations: typing.List[str] = [], labels: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_LABELS', annotations: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_ANNOTATIONS', use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = '/metaflow_temp', persistent_volume_claims: typing.Optional[typing.Dict[str, str]] = None, shared_memory: typing.Optional[int] = None, port: typing.Optional[int] = None, compute_pool: typing.Optional[str] = None, hostname_resolution_timeout: int = 600, qos: str = 'Burstable', 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]]]:
425
407
  """
426
- This decorator is used to run Ollama APIs as Metaflow task sidecars.
427
-
428
- User code call
429
- --------------
430
- @ollama(
431
- models=[...],
432
- ...
433
- )
434
-
435
- Valid backend options
436
- ---------------------
437
- - 'local': Run as a separate process on the local task machine.
438
- - (TODO) 'managed': Outerbounds hosts and selects compute provider.
439
- - (TODO) 'remote': Spin up separate instance to serve Ollama models.
440
-
441
- Valid model options
442
- -------------------
443
- Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
444
-
445
-
446
- Parameters
447
- ----------
448
- models: list[str]
449
- List of Ollama containers running models in sidecars.
450
- backend: str
451
- Determines where and how to run the Ollama process.
452
- force_pull: bool
453
- Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
454
- skip_push_check: bool
455
- Whether to skip the check that populates/overwrites remote cache on terminating an ollama model.
456
- debug: bool
457
- Whether to turn on verbose debugging logs.
458
- """
459
- ...
460
-
461
- @typing.overload
462
- 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]]]:
463
- """
464
- Specifies environment variables to be set prior to the execution of a step.
465
-
466
-
467
- Parameters
468
- ----------
469
- vars : Dict[str, str], default {}
470
- Dictionary of environment variables to set.
471
- """
472
- ...
473
-
474
- @typing.overload
475
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
476
- ...
477
-
478
- @typing.overload
479
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
480
- ...
481
-
482
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
483
- """
484
- Specifies environment variables to be set prior to the execution of a step.
485
-
486
-
487
- Parameters
488
- ----------
489
- vars : Dict[str, str], default {}
490
- Dictionary of environment variables to set.
491
- """
492
- ...
493
-
494
- @typing.overload
495
- 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]]]:
496
- """
497
- Specifies the PyPI packages for the step.
498
-
499
- Information in this decorator will augment any
500
- attributes set in the `@pyi_base` flow-level decorator. Hence,
501
- you can use `@pypi_base` to set packages required by all
502
- steps and use `@pypi` to specify step-specific overrides.
408
+ Specifies that this step should execute on Kubernetes.
503
409
 
504
410
 
505
411
  Parameters
506
412
  ----------
507
- packages : Dict[str, str], default: {}
508
- Packages to use for this step. The key is the name of the package
509
- and the value is the version to use.
510
- python : str, optional, default: None
511
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
512
- that the version used will correspond to the version of the Python interpreter used to start the run.
513
- """
514
- ...
515
-
516
- @typing.overload
517
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
518
- ...
519
-
520
- @typing.overload
521
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
522
- ...
523
-
524
- 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):
525
- """
526
- Specifies the PyPI packages for the step.
527
-
528
- Information in this decorator will augment any
529
- attributes set in the `@pyi_base` flow-level decorator. Hence,
530
- you can use `@pypi_base` to set packages required by all
531
- steps and use `@pypi` to specify step-specific overrides.
532
-
413
+ cpu : int, default 1
414
+ Number of CPUs required for this step. If `@resources` is
415
+ also present, the maximum value from all decorators is used.
416
+ memory : int, default 4096
417
+ Memory size (in MB) required for this step. If
418
+ `@resources` is also present, the maximum value from all decorators is
419
+ used.
420
+ disk : int, default 10240
421
+ Disk size (in MB) required for this step. If
422
+ `@resources` is also present, the maximum value from all decorators is
423
+ used.
424
+ image : str, optional, default None
425
+ Docker image to use when launching on Kubernetes. If not specified, and
426
+ METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
427
+ not, a default Docker image mapping to the current version of Python is used.
428
+ image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
429
+ If given, the imagePullPolicy to be applied to the Docker image of the step.
430
+ service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
431
+ Kubernetes service account to use when launching pod in Kubernetes.
432
+ secrets : List[str], optional, default None
433
+ Kubernetes secrets to use when launching pod in Kubernetes. These
434
+ secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
435
+ in Metaflow configuration.
436
+ node_selector: Union[Dict[str,str], str], optional, default None
437
+ Kubernetes node selector(s) to apply to the pod running the task.
438
+ Can be passed in as a comma separated string of values e.g.
439
+ 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
440
+ {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
441
+ namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
442
+ Kubernetes namespace to use when launching pod in Kubernetes.
443
+ gpu : int, optional, default None
444
+ Number of GPUs required for this step. A value of zero implies that
445
+ the scheduled node should not have GPUs.
446
+ gpu_vendor : str, default KUBERNETES_GPU_VENDOR
447
+ The vendor of the GPUs to be used for this step.
448
+ tolerations : List[str], default []
449
+ The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
450
+ Kubernetes tolerations to use when launching pod in Kubernetes.
451
+ labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
452
+ Kubernetes labels to use when launching pod in Kubernetes.
453
+ annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
454
+ Kubernetes annotations to use when launching pod in Kubernetes.
455
+ use_tmpfs : bool, default False
456
+ This enables an explicit tmpfs mount for this step.
457
+ tmpfs_tempdir : bool, default True
458
+ sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
459
+ tmpfs_size : int, optional, default: None
460
+ The value for the size (in MiB) of the tmpfs mount for this step.
461
+ This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
462
+ memory allocated for this step.
463
+ tmpfs_path : str, optional, default /metaflow_temp
464
+ Path to tmpfs mount for this step.
465
+ persistent_volume_claims : Dict[str, str], optional, default None
466
+ A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
467
+ volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
468
+ shared_memory: int, optional
469
+ Shared memory size (in MiB) required for this step
470
+ port: int, optional
471
+ Port number to specify in the Kubernetes job object
472
+ compute_pool : str, optional, default None
473
+ Compute pool to be used for for this step.
474
+ If not specified, any accessible compute pool within the perimeter is used.
475
+ hostname_resolution_timeout: int, default 10 * 60
476
+ Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
477
+ Only applicable when @parallel is used.
478
+ qos: str, default: Burstable
479
+ Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
533
480
 
534
- Parameters
535
- ----------
536
- packages : Dict[str, str], default: {}
537
- Packages to use for this step. The key is the name of the package
538
- and the value is the version to use.
539
- python : str, optional, default: None
540
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
541
- that the version used will correspond to the version of the Python interpreter used to start the run.
481
+ security_context: Dict[str, Any], optional, default None
482
+ Container security context. Applies to the task container. Allows the following keys:
483
+ - privileged: bool, optional, default None
484
+ - allow_privilege_escalation: bool, optional, default None
485
+ - run_as_user: int, optional, default None
486
+ - run_as_group: int, optional, default None
487
+ - run_as_non_root: bool, optional, default None
542
488
  """
543
489
  ...
544
490
 
@@ -594,131 +540,95 @@ def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
594
540
  ...
595
541
 
596
542
  @typing.overload
597
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
598
- """
599
- Decorator prototype for all step decorators. This function gets specialized
600
- and imported for all decorators types by _import_plugin_decorators().
601
- """
602
- ...
603
-
604
- @typing.overload
605
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
606
- ...
607
-
608
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
609
- """
610
- Decorator prototype for all step decorators. This function gets specialized
611
- and imported for all decorators types by _import_plugin_decorators().
612
- """
613
- ...
614
-
615
- @typing.overload
616
- 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]]]:
543
+ def checkpoint(*, load_policy: str = 'fresh', temp_dir_root: str = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
617
544
  """
618
- Creates a human-readable report, a Metaflow Card, after this step completes.
545
+ Enables checkpointing for a step.
619
546
 
620
- Note that you may add multiple `@card` decorators in a step with different parameters.
621
547
 
622
548
 
623
549
  Parameters
624
550
  ----------
625
- type : str, default 'default'
626
- Card type.
627
- id : str, optional, default None
628
- If multiple cards are present, use this id to identify this card.
629
- options : Dict[str, Any], default {}
630
- Options passed to the card. The contents depend on the card type.
631
- timeout : int, default 45
632
- Interrupt reporting if it takes more than this many seconds.
551
+ load_policy : str, default: "fresh"
552
+ The policy for loading the checkpoint. The following policies are supported:
553
+ - "eager": Loads the the latest available checkpoint within the namespace.
554
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
555
+ will be loaded at the start of the task.
556
+ - "none": Do not load any checkpoint
557
+ - "fresh": Loads the lastest checkpoint created within the running Task.
558
+ This mode helps loading checkpoints across various retry attempts of the same task.
559
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
560
+ created within the task will be loaded when the task is retries execution on failure.
561
+
562
+ temp_dir_root : str, default: None
563
+ The root directory under which `current.checkpoint.directory` will be created.
633
564
  """
634
565
  ...
635
566
 
636
567
  @typing.overload
637
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
568
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
638
569
  ...
639
570
 
640
571
  @typing.overload
641
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
572
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
642
573
  ...
643
574
 
644
- 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):
575
+ def checkpoint(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, load_policy: str = 'fresh', temp_dir_root: str = None):
645
576
  """
646
- Creates a human-readable report, a Metaflow Card, after this step completes.
577
+ Enables checkpointing for a step.
647
578
 
648
- Note that you may add multiple `@card` decorators in a step with different parameters.
649
579
 
650
580
 
651
581
  Parameters
652
582
  ----------
653
- type : str, default 'default'
654
- Card type.
655
- id : str, optional, default None
656
- If multiple cards are present, use this id to identify this card.
657
- options : Dict[str, Any], default {}
658
- Options passed to the card. The contents depend on the card type.
659
- timeout : int, default 45
660
- Interrupt reporting if it takes more than this many seconds.
583
+ load_policy : str, default: "fresh"
584
+ The policy for loading the checkpoint. The following policies are supported:
585
+ - "eager": Loads the the latest available checkpoint within the namespace.
586
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
587
+ will be loaded at the start of the task.
588
+ - "none": Do not load any checkpoint
589
+ - "fresh": Loads the lastest checkpoint created within the running Task.
590
+ This mode helps loading checkpoints across various retry attempts of the same task.
591
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
592
+ created within the task will be loaded when the task is retries execution on failure.
593
+
594
+ temp_dir_root : str, default: None
595
+ The root directory under which `current.checkpoint.directory` will be created.
661
596
  """
662
597
  ...
663
598
 
664
- @typing.overload
665
- 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]]]:
599
+ 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]]]:
666
600
  """
667
- Enables loading / saving of models within a step.
668
-
601
+ Specifies that this step is used to deploy an instance of the app.
602
+ Requires that self.app_name, self.app_port, self.entrypoint and self.deployDir is set.
669
603
 
670
604
 
671
605
  Parameters
672
606
  ----------
673
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
674
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
675
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
676
- - `current.checkpoint`
677
- - `current.model`
678
- - `current.huggingface_hub`
679
-
680
- 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
681
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
682
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
683
-
684
- temp_dir_root : str, default: None
685
- The root directory under which `current.model.loaded` will store loaded models
607
+ app_port : int
608
+ Number of GPUs to use.
609
+ app_name : str
610
+ Name of the app to deploy.
686
611
  """
687
612
  ...
688
613
 
689
614
  @typing.overload
690
- def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
615
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
616
+ """
617
+ Internal decorator to support Fast bakery
618
+ """
691
619
  ...
692
620
 
693
621
  @typing.overload
694
- def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
622
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
695
623
  ...
696
624
 
697
- 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):
625
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
698
626
  """
699
- Enables loading / saving of models within a step.
700
-
701
-
702
-
703
- Parameters
704
- ----------
705
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
706
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
707
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
708
- - `current.checkpoint`
709
- - `current.model`
710
- - `current.huggingface_hub`
711
-
712
- 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
713
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
714
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
715
-
716
- temp_dir_root : str, default: None
717
- The root directory under which `current.model.loaded` will store loaded models
627
+ Internal decorator to support Fast bakery
718
628
  """
719
629
  ...
720
630
 
721
- 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]]]:
631
+ 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]]]:
722
632
  """
723
633
  Specifies that this step should execute on DGX cloud.
724
634
 
@@ -729,454 +639,330 @@ def nvct(*, gpu: int, gpu_type: str) -> typing.Callable[[typing.Union[typing.Cal
729
639
  Number of GPUs to use.
730
640
  gpu_type : str
731
641
  Type of Nvidia GPU to use.
642
+ queue_timeout : int
643
+ Time to keep the job in NVCF's queue.
644
+ """
645
+ ...
646
+
647
+ 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]]]:
648
+ """
649
+ Decorator that helps cache, version and store models/datasets from huggingface hub.
650
+
651
+
652
+ Parameters
653
+ ----------
654
+ temp_dir_root : str, optional
655
+ The root directory that will hold the temporary directory where objects will be downloaded.
656
+
657
+ load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
658
+ The list of repos (models/datasets) to load.
659
+
660
+ Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
661
+
662
+ - If repo (model/dataset) is not found in the datastore:
663
+ - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
664
+ - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
665
+ - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
666
+
667
+ - If repo is found in the datastore:
668
+ - Loads it directly from datastore to local path (can be temporary directory or specified path)
732
669
  """
733
670
  ...
734
671
 
735
672
  @typing.overload
736
- 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]]]:
673
+ 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]]]:
737
674
  """
738
- Specifies the Conda environment for the step.
675
+ Specifies the PyPI packages for the step.
739
676
 
740
677
  Information in this decorator will augment any
741
- attributes set in the `@conda_base` flow-level decorator. Hence,
742
- you can use `@conda_base` to set packages required by all
743
- steps and use `@conda` to specify step-specific overrides.
678
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
679
+ you can use `@pypi_base` to set packages required by all
680
+ steps and use `@pypi` to specify step-specific overrides.
744
681
 
745
682
 
746
683
  Parameters
747
684
  ----------
748
- packages : Dict[str, str], default {}
685
+ packages : Dict[str, str], default: {}
749
686
  Packages to use for this step. The key is the name of the package
750
687
  and the value is the version to use.
751
- libraries : Dict[str, str], default {}
752
- Supported for backward compatibility. When used with packages, packages will take precedence.
753
- python : str, optional, default None
688
+ python : str, optional, default: None
754
689
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
755
690
  that the version used will correspond to the version of the Python interpreter used to start the run.
756
- disabled : bool, default False
757
- If set to True, disables @conda.
758
691
  """
759
692
  ...
760
693
 
761
694
  @typing.overload
762
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
695
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
763
696
  ...
764
697
 
765
698
  @typing.overload
766
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
699
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
767
700
  ...
768
701
 
769
- 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):
702
+ 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):
770
703
  """
771
- Specifies the Conda environment for the step.
704
+ Specifies the PyPI packages for the step.
772
705
 
773
706
  Information in this decorator will augment any
774
- attributes set in the `@conda_base` flow-level decorator. Hence,
775
- you can use `@conda_base` to set packages required by all
776
- steps and use `@conda` to specify step-specific overrides.
707
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
708
+ you can use `@pypi_base` to set packages required by all
709
+ steps and use `@pypi` to specify step-specific overrides.
777
710
 
778
711
 
779
712
  Parameters
780
713
  ----------
781
- packages : Dict[str, str], default {}
714
+ packages : Dict[str, str], default: {}
782
715
  Packages to use for this step. The key is the name of the package
783
716
  and the value is the version to use.
784
- libraries : Dict[str, str], default {}
785
- Supported for backward compatibility. When used with packages, packages will take precedence.
786
- python : str, optional, default None
717
+ python : str, optional, default: None
787
718
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
788
719
  that the version used will correspond to the version of the Python interpreter used to start the run.
789
- disabled : bool, default False
790
- If set to True, disables @conda.
791
720
  """
792
721
  ...
793
722
 
794
- def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: typing.Optional[str] = None, image_pull_policy: str = 'KUBERNETES_IMAGE_PULL_POLICY', service_account: str = 'METAFLOW_KUBERNETES_SERVICE_ACCOUNT', secrets: typing.Optional[typing.List[str]] = None, node_selector: typing.Union[typing.Dict[str, str], str, None] = None, namespace: str = 'METAFLOW_KUBERNETES_NAMESPACE', gpu: typing.Optional[int] = None, gpu_vendor: str = 'KUBERNETES_GPU_VENDOR', tolerations: typing.List[str] = [], labels: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_LABELS', annotations: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_ANNOTATIONS', use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = '/metaflow_temp', persistent_volume_claims: typing.Optional[typing.Dict[str, str]] = None, shared_memory: typing.Optional[int] = None, port: typing.Optional[int] = None, compute_pool: typing.Optional[str] = None, hostname_resolution_timeout: int = 600, qos: str = 'Burstable', 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]]]:
723
+ def ollama(*, models: list, backend: str, force_pull: bool, skip_push_check: bool, debug: bool) -> 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]]]:
795
724
  """
796
- Specifies that this step should execute on Kubernetes.
725
+ This decorator is used to run Ollama APIs as Metaflow task sidecars.
726
+
727
+ User code call
728
+ --------------
729
+ @ollama(
730
+ models=[...],
731
+ ...
732
+ )
733
+
734
+ Valid backend options
735
+ ---------------------
736
+ - 'local': Run as a separate process on the local task machine.
737
+ - (TODO) 'managed': Outerbounds hosts and selects compute provider.
738
+ - (TODO) 'remote': Spin up separate instance to serve Ollama models.
739
+
740
+ Valid model options
741
+ -------------------
742
+ Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
797
743
 
798
744
 
799
745
  Parameters
800
746
  ----------
801
- cpu : int, default 1
802
- Number of CPUs required for this step. If `@resources` is
803
- also present, the maximum value from all decorators is used.
804
- memory : int, default 4096
805
- Memory size (in MB) required for this step. If
806
- `@resources` is also present, the maximum value from all decorators is
807
- used.
808
- disk : int, default 10240
809
- Disk size (in MB) required for this step. If
810
- `@resources` is also present, the maximum value from all decorators is
811
- used.
812
- image : str, optional, default None
813
- Docker image to use when launching on Kubernetes. If not specified, and
814
- METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
815
- not, a default Docker image mapping to the current version of Python is used.
816
- image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
817
- If given, the imagePullPolicy to be applied to the Docker image of the step.
818
- service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
819
- Kubernetes service account to use when launching pod in Kubernetes.
820
- secrets : List[str], optional, default None
821
- Kubernetes secrets to use when launching pod in Kubernetes. These
822
- secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
823
- in Metaflow configuration.
824
- node_selector: Union[Dict[str,str], str], optional, default None
825
- Kubernetes node selector(s) to apply to the pod running the task.
826
- Can be passed in as a comma separated string of values e.g.
827
- 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
828
- {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
829
- namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
830
- Kubernetes namespace to use when launching pod in Kubernetes.
831
- gpu : int, optional, default None
832
- Number of GPUs required for this step. A value of zero implies that
833
- the scheduled node should not have GPUs.
834
- gpu_vendor : str, default KUBERNETES_GPU_VENDOR
835
- The vendor of the GPUs to be used for this step.
836
- tolerations : List[str], default []
837
- The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
838
- Kubernetes tolerations to use when launching pod in Kubernetes.
839
- labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
840
- Kubernetes labels to use when launching pod in Kubernetes.
841
- annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
842
- Kubernetes annotations to use when launching pod in Kubernetes.
843
- use_tmpfs : bool, default False
844
- This enables an explicit tmpfs mount for this step.
845
- tmpfs_tempdir : bool, default True
846
- sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
847
- tmpfs_size : int, optional, default: None
848
- The value for the size (in MiB) of the tmpfs mount for this step.
849
- This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
850
- memory allocated for this step.
851
- tmpfs_path : str, optional, default /metaflow_temp
852
- Path to tmpfs mount for this step.
853
- persistent_volume_claims : Dict[str, str], optional, default None
854
- A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
855
- volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
856
- shared_memory: int, optional
857
- Shared memory size (in MiB) required for this step
858
- port: int, optional
859
- Port number to specify in the Kubernetes job object
860
- compute_pool : str, optional, default None
861
- Compute pool to be used for for this step.
862
- If not specified, any accessible compute pool within the perimeter is used.
863
- hostname_resolution_timeout: int, default 10 * 60
864
- Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
865
- Only applicable when @parallel is used.
866
- qos: str, default: Burstable
867
- Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
747
+ models: list[str]
748
+ List of Ollama containers running models in sidecars.
749
+ backend: str
750
+ Determines where and how to run the Ollama process.
751
+ force_pull: bool
752
+ Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
753
+ skip_push_check: bool
754
+ Whether to skip the check that populates/overwrites remote cache on terminating an ollama model.
755
+ debug: bool
756
+ Whether to turn on verbose debugging logs.
757
+ """
758
+ ...
759
+
760
+ 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]]]:
761
+ """
762
+ Specifies that this step should execute on DGX cloud.
868
763
 
869
- security_context: Dict[str, Any], optional, default None
870
- Container security context. Applies to the task container. Allows the following keys:
871
- - privileged: bool, optional, default None
872
- - allow_privilege_escalation: bool, optional, default None
873
- - run_as_user: int, optional, default None
874
- - run_as_group: int, optional, default None
875
- - run_as_non_root: bool, optional, default None
764
+
765
+ Parameters
766
+ ----------
767
+ gpu : int
768
+ Number of GPUs to use.
769
+ gpu_type : str
770
+ Type of Nvidia GPU to use.
876
771
  """
877
772
  ...
878
773
 
879
774
  @typing.overload
880
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
775
+ 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]]]:
881
776
  """
882
- Internal decorator to support Fast bakery
777
+ Specifies secrets to be retrieved and injected as environment variables prior to
778
+ the execution of a step.
779
+
780
+
781
+ Parameters
782
+ ----------
783
+ sources : List[Union[str, Dict[str, Any]]], default: []
784
+ List of secret specs, defining how the secrets are to be retrieved
883
785
  """
884
786
  ...
885
787
 
886
788
  @typing.overload
887
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
789
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
888
790
  ...
889
791
 
890
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
891
- """
892
- Internal decorator to support Fast bakery
893
- """
792
+ @typing.overload
793
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
894
794
  ...
895
795
 
896
- @typing.overload
897
- def checkpoint(*, load_policy: str = 'fresh', temp_dir_root: str = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
796
+ 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]]] = []):
898
797
  """
899
- Enables checkpointing for a step.
900
-
798
+ Specifies secrets to be retrieved and injected as environment variables prior to
799
+ the execution of a step.
901
800
 
902
801
 
903
802
  Parameters
904
803
  ----------
905
- load_policy : str, default: "fresh"
906
- The policy for loading the checkpoint. The following policies are supported:
907
- - "eager": Loads the the latest available checkpoint within the namespace.
908
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
909
- will be loaded at the start of the task.
910
- - "none": Do not load any checkpoint
911
- - "fresh": Loads the lastest checkpoint created within the running Task.
912
- This mode helps loading checkpoints across various retry attempts of the same task.
913
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
914
- created within the task will be loaded when the task is retries execution on failure.
915
-
916
- temp_dir_root : str, default: None
917
- The root directory under which `current.checkpoint.directory` will be created.
804
+ sources : List[Union[str, Dict[str, Any]]], default: []
805
+ List of secret specs, defining how the secrets are to be retrieved
918
806
  """
919
807
  ...
920
808
 
921
809
  @typing.overload
922
- def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
810
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
811
+ """
812
+ Decorator prototype for all step decorators. This function gets specialized
813
+ and imported for all decorators types by _import_plugin_decorators().
814
+ """
923
815
  ...
924
816
 
925
817
  @typing.overload
926
- def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
818
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
927
819
  ...
928
820
 
929
- def checkpoint(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, load_policy: str = 'fresh', temp_dir_root: str = None):
821
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
930
822
  """
931
- Enables checkpointing for a step.
932
-
933
-
934
-
935
- Parameters
936
- ----------
937
- load_policy : str, default: "fresh"
938
- The policy for loading the checkpoint. The following policies are supported:
939
- - "eager": Loads the the latest available checkpoint within the namespace.
940
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
941
- will be loaded at the start of the task.
942
- - "none": Do not load any checkpoint
943
- - "fresh": Loads the lastest checkpoint created within the running Task.
944
- This mode helps loading checkpoints across various retry attempts of the same task.
945
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
946
- created within the task will be loaded when the task is retries execution on failure.
947
-
948
- temp_dir_root : str, default: None
949
- The root directory under which `current.checkpoint.directory` will be created.
823
+ Decorator prototype for all step decorators. This function gets specialized
824
+ and imported for all decorators types by _import_plugin_decorators().
950
825
  """
951
826
  ...
952
827
 
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]]]:
828
+ @typing.overload
829
+ 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]]]:
954
830
  """
955
- Specifies that this step should execute on DGX cloud.
831
+ Creates a human-readable report, a Metaflow Card, after this step completes.
832
+
833
+ Note that you may add multiple `@card` decorators in a step with different parameters.
956
834
 
957
835
 
958
836
  Parameters
959
837
  ----------
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.
838
+ type : str, default 'default'
839
+ Card type.
840
+ id : str, optional, default None
841
+ If multiple cards are present, use this id to identify this card.
842
+ options : Dict[str, Any], default {}
843
+ Options passed to the card. The contents depend on the card type.
844
+ timeout : int, default 45
845
+ Interrupt reporting if it takes more than this many seconds.
966
846
  """
967
847
  ...
968
848
 
969
- 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]]:
849
+ @typing.overload
850
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
851
+ ...
852
+
853
+ @typing.overload
854
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
855
+ ...
856
+
857
+ 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):
970
858
  """
971
- 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.
972
- 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.
859
+ Creates a human-readable report, a Metaflow Card, after this step completes.
860
+
861
+ Note that you may add multiple `@card` decorators in a step with different parameters.
973
862
 
974
863
 
975
864
  Parameters
976
865
  ----------
977
- timeout : int
978
- Time, in seconds before the task times out and fails. (Default: 3600)
979
- poke_interval : int
980
- Time in seconds that the job should wait in between each try. (Default: 60)
981
- mode : str
982
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
983
- exponential_backoff : bool
984
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
985
- pool : str
986
- the slot pool this task should run in,
987
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
988
- soft_fail : bool
989
- Set to true to mark the task as SKIPPED on failure. (Default: False)
990
- name : str
991
- Name of the sensor on Airflow
992
- description : str
993
- Description of sensor in the Airflow UI
994
- external_dag_id : str
995
- The dag_id that contains the task you want to wait for.
996
- external_task_ids : List[str]
997
- The list of task_ids that you want to wait for.
998
- If None (default value) the sensor waits for the DAG. (Default: None)
999
- allowed_states : List[str]
1000
- Iterable of allowed states, (Default: ['success'])
1001
- failed_states : List[str]
1002
- Iterable of failed or dis-allowed states. (Default: None)
1003
- execution_delta : datetime.timedelta
1004
- time difference with the previous execution to look at,
1005
- the default is the same logical date as the current task or DAG. (Default: None)
1006
- check_existence: bool
1007
- Set to True to check if the external task exists or check if
1008
- the DAG to wait for exists. (Default: True)
866
+ type : str, default 'default'
867
+ Card type.
868
+ id : str, optional, default None
869
+ If multiple cards are present, use this id to identify this card.
870
+ options : Dict[str, Any], default {}
871
+ Options passed to the card. The contents depend on the card type.
872
+ timeout : int, default 45
873
+ Interrupt reporting if it takes more than this many seconds.
1009
874
  """
1010
875
  ...
1011
876
 
1012
- def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
877
+ @typing.overload
878
+ 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]]]:
1013
879
  """
1014
- Specifies what flows belong to the same project.
1015
-
1016
- A project-specific namespace is created for all flows that
1017
- use the same `@project(name)`.
880
+ Specifies environment variables to be set prior to the execution of a step.
1018
881
 
1019
882
 
1020
883
  Parameters
1021
884
  ----------
1022
- name : str
1023
- Project name. Make sure that the name is unique amongst all
1024
- projects that use the same production scheduler. The name may
1025
- contain only lowercase alphanumeric characters and underscores.
1026
-
1027
- branch : Optional[str], default None
1028
- The branch to use. If not specified, the branch is set to
1029
- `user.<username>` unless `production` is set to `True`. This can
1030
- also be set on the command line using `--branch` as a top-level option.
1031
- It is an error to specify `branch` in the decorator and on the command line.
1032
-
1033
- production : bool, default False
1034
- Whether or not the branch is the production branch. This can also be set on the
1035
- command line using `--production` as a top-level option. It is an error to specify
1036
- `production` in the decorator and on the command line.
1037
- The project branch name will be:
1038
- - if `branch` is specified:
1039
- - if `production` is True: `prod.<branch>`
1040
- - if `production` is False: `test.<branch>`
1041
- - if `branch` is not specified:
1042
- - if `production` is True: `prod`
1043
- - if `production` is False: `user.<username>`
885
+ vars : Dict[str, str], default {}
886
+ Dictionary of environment variables to set.
1044
887
  """
1045
888
  ...
1046
889
 
1047
- 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]]:
890
+ @typing.overload
891
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
892
+ ...
893
+
894
+ @typing.overload
895
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
896
+ ...
897
+
898
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
1048
899
  """
1049
- 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)
1050
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1051
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1052
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
1053
- starts only after all sensors finish.
900
+ Specifies environment variables to be set prior to the execution of a step.
1054
901
 
1055
902
 
1056
903
  Parameters
1057
904
  ----------
1058
- timeout : int
1059
- Time, in seconds before the task times out and fails. (Default: 3600)
1060
- poke_interval : int
1061
- Time in seconds that the job should wait in between each try. (Default: 60)
1062
- mode : str
1063
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1064
- exponential_backoff : bool
1065
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1066
- pool : str
1067
- the slot pool this task should run in,
1068
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1069
- soft_fail : bool
1070
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1071
- name : str
1072
- Name of the sensor on Airflow
1073
- description : str
1074
- Description of sensor in the Airflow UI
1075
- bucket_key : Union[str, List[str]]
1076
- The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1077
- When it's specified as a full s3:// url, please leave `bucket_name` as None
1078
- bucket_name : str
1079
- Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1080
- When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1081
- wildcard_match : bool
1082
- whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1083
- aws_conn_id : str
1084
- a reference to the s3 connection on Airflow. (Default: None)
1085
- verify : bool
1086
- Whether or not to verify SSL certificates for S3 connection. (Default: None)
905
+ vars : Dict[str, str], default {}
906
+ Dictionary of environment variables to set.
1087
907
  """
1088
908
  ...
1089
909
 
1090
910
  @typing.overload
1091
- 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]]:
911
+ def timeout(*, seconds: int = 0, minutes: int = 0, hours: int = 0) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
1092
912
  """
1093
- Specifies the event(s) that this flow depends on.
913
+ Specifies a timeout for your step.
1094
914
 
1095
- ```
1096
- @trigger(event='foo')
1097
- ```
1098
- or
1099
- ```
1100
- @trigger(events=['foo', 'bar'])
1101
- ```
915
+ This decorator is useful if this step may hang indefinitely.
1102
916
 
1103
- Additionally, you can specify the parameter mappings
1104
- to map event payload to Metaflow parameters for the flow.
1105
- ```
1106
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1107
- ```
1108
- or
1109
- ```
1110
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1111
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1112
- ```
917
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
918
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
919
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
1113
920
 
1114
- 'parameters' can also be a list of strings and tuples like so:
1115
- ```
1116
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1117
- ```
1118
- This is equivalent to:
1119
- ```
1120
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1121
- ```
921
+ Note that all the values specified in parameters are added together so if you specify
922
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
1122
923
 
1123
924
 
1124
925
  Parameters
1125
926
  ----------
1126
- event : Union[str, Dict[str, Any]], optional, default None
1127
- Event dependency for this flow.
1128
- events : List[Union[str, Dict[str, Any]]], default []
1129
- Events dependency for this flow.
1130
- options : Dict[str, Any], default {}
1131
- Backend-specific configuration for tuning eventing behavior.
927
+ seconds : int, default 0
928
+ Number of seconds to wait prior to timing out.
929
+ minutes : int, default 0
930
+ Number of minutes to wait prior to timing out.
931
+ hours : int, default 0
932
+ Number of hours to wait prior to timing out.
1132
933
  """
1133
934
  ...
1134
935
 
1135
936
  @typing.overload
1136
- def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
937
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1137
938
  ...
1138
939
 
1139
- 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] = {}):
1140
- """
1141
- Specifies the event(s) that this flow depends on.
1142
-
1143
- ```
1144
- @trigger(event='foo')
1145
- ```
1146
- or
1147
- ```
1148
- @trigger(events=['foo', 'bar'])
1149
- ```
940
+ @typing.overload
941
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
942
+ ...
943
+
944
+ def timeout(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, seconds: int = 0, minutes: int = 0, hours: int = 0):
945
+ """
946
+ Specifies a timeout for your step.
1150
947
 
1151
- Additionally, you can specify the parameter mappings
1152
- to map event payload to Metaflow parameters for the flow.
1153
- ```
1154
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1155
- ```
1156
- or
1157
- ```
1158
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1159
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1160
- ```
948
+ This decorator is useful if this step may hang indefinitely.
1161
949
 
1162
- 'parameters' can also be a list of strings and tuples like so:
1163
- ```
1164
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1165
- ```
1166
- This is equivalent to:
1167
- ```
1168
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1169
- ```
950
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
951
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
952
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
953
+
954
+ Note that all the values specified in parameters are added together so if you specify
955
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
1170
956
 
1171
957
 
1172
958
  Parameters
1173
959
  ----------
1174
- event : Union[str, Dict[str, Any]], optional, default None
1175
- Event dependency for this flow.
1176
- events : List[Union[str, Dict[str, Any]]], default []
1177
- Events dependency for this flow.
1178
- options : Dict[str, Any], default {}
1179
- Backend-specific configuration for tuning eventing behavior.
960
+ seconds : int, default 0
961
+ Number of seconds to wait prior to timing out.
962
+ minutes : int, default 0
963
+ Number of minutes to wait prior to timing out.
964
+ hours : int, default 0
965
+ Number of hours to wait prior to timing out.
1180
966
  """
1181
967
  ...
1182
968
 
@@ -1345,6 +1131,127 @@ def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None)
1345
1131
  """
1346
1132
  ...
1347
1133
 
1134
+ def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1135
+ """
1136
+ Specifies what flows belong to the same project.
1137
+
1138
+ A project-specific namespace is created for all flows that
1139
+ use the same `@project(name)`.
1140
+
1141
+
1142
+ Parameters
1143
+ ----------
1144
+ name : str
1145
+ Project name. Make sure that the name is unique amongst all
1146
+ projects that use the same production scheduler. The name may
1147
+ contain only lowercase alphanumeric characters and underscores.
1148
+
1149
+ branch : Optional[str], default None
1150
+ The branch to use. If not specified, the branch is set to
1151
+ `user.<username>` unless `production` is set to `True`. This can
1152
+ also be set on the command line using `--branch` as a top-level option.
1153
+ It is an error to specify `branch` in the decorator and on the command line.
1154
+
1155
+ production : bool, default False
1156
+ Whether or not the branch is the production branch. This can also be set on the
1157
+ command line using `--production` as a top-level option. It is an error to specify
1158
+ `production` in the decorator and on the command line.
1159
+ The project branch name will be:
1160
+ - if `branch` is specified:
1161
+ - if `production` is True: `prod.<branch>`
1162
+ - if `production` is False: `test.<branch>`
1163
+ - if `branch` is not specified:
1164
+ - if `production` is True: `prod`
1165
+ - if `production` is False: `user.<username>`
1166
+ """
1167
+ ...
1168
+
1169
+ 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]]:
1170
+ """
1171
+ 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)
1172
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1173
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1174
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
1175
+ starts only after all sensors finish.
1176
+
1177
+
1178
+ Parameters
1179
+ ----------
1180
+ timeout : int
1181
+ Time, in seconds before the task times out and fails. (Default: 3600)
1182
+ poke_interval : int
1183
+ Time in seconds that the job should wait in between each try. (Default: 60)
1184
+ mode : str
1185
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1186
+ exponential_backoff : bool
1187
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1188
+ pool : str
1189
+ the slot pool this task should run in,
1190
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1191
+ soft_fail : bool
1192
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1193
+ name : str
1194
+ Name of the sensor on Airflow
1195
+ description : str
1196
+ Description of sensor in the Airflow UI
1197
+ bucket_key : Union[str, List[str]]
1198
+ The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1199
+ When it's specified as a full s3:// url, please leave `bucket_name` as None
1200
+ bucket_name : str
1201
+ Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1202
+ When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1203
+ wildcard_match : bool
1204
+ whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1205
+ aws_conn_id : str
1206
+ a reference to the s3 connection on Airflow. (Default: None)
1207
+ verify : bool
1208
+ Whether or not to verify SSL certificates for S3 connection. (Default: None)
1209
+ """
1210
+ ...
1211
+
1212
+ 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]]:
1213
+ """
1214
+ 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.
1215
+ 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.
1216
+
1217
+
1218
+ Parameters
1219
+ ----------
1220
+ timeout : int
1221
+ Time, in seconds before the task times out and fails. (Default: 3600)
1222
+ poke_interval : int
1223
+ Time in seconds that the job should wait in between each try. (Default: 60)
1224
+ mode : str
1225
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1226
+ exponential_backoff : bool
1227
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1228
+ pool : str
1229
+ the slot pool this task should run in,
1230
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1231
+ soft_fail : bool
1232
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1233
+ name : str
1234
+ Name of the sensor on Airflow
1235
+ description : str
1236
+ Description of sensor in the Airflow UI
1237
+ external_dag_id : str
1238
+ The dag_id that contains the task you want to wait for.
1239
+ external_task_ids : List[str]
1240
+ The list of task_ids that you want to wait for.
1241
+ If None (default value) the sensor waits for the DAG. (Default: None)
1242
+ allowed_states : List[str]
1243
+ Iterable of allowed states, (Default: ['success'])
1244
+ failed_states : List[str]
1245
+ Iterable of failed or dis-allowed states. (Default: None)
1246
+ execution_delta : datetime.timedelta
1247
+ time difference with the previous execution to look at,
1248
+ the default is the same logical date as the current task or DAG. (Default: None)
1249
+ check_existence: bool
1250
+ Set to True to check if the external task exists or check if
1251
+ the DAG to wait for exists. (Default: True)
1252
+ """
1253
+ ...
1254
+
1348
1255
  @typing.overload
1349
1256
  def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1350
1257
  """
@@ -1487,6 +1394,99 @@ def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *
1487
1394
  """
1488
1395
  ...
1489
1396
 
1397
+ @typing.overload
1398
+ 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]]:
1399
+ """
1400
+ Specifies the event(s) that this flow depends on.
1401
+
1402
+ ```
1403
+ @trigger(event='foo')
1404
+ ```
1405
+ or
1406
+ ```
1407
+ @trigger(events=['foo', 'bar'])
1408
+ ```
1409
+
1410
+ Additionally, you can specify the parameter mappings
1411
+ to map event payload to Metaflow parameters for the flow.
1412
+ ```
1413
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1414
+ ```
1415
+ or
1416
+ ```
1417
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1418
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1419
+ ```
1420
+
1421
+ 'parameters' can also be a list of strings and tuples like so:
1422
+ ```
1423
+ @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1424
+ ```
1425
+ This is equivalent to:
1426
+ ```
1427
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1428
+ ```
1429
+
1430
+
1431
+ Parameters
1432
+ ----------
1433
+ event : Union[str, Dict[str, Any]], optional, default None
1434
+ Event dependency for this flow.
1435
+ events : List[Union[str, Dict[str, Any]]], default []
1436
+ Events dependency for this flow.
1437
+ options : Dict[str, Any], default {}
1438
+ Backend-specific configuration for tuning eventing behavior.
1439
+ """
1440
+ ...
1441
+
1442
+ @typing.overload
1443
+ def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1444
+ ...
1445
+
1446
+ 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] = {}):
1447
+ """
1448
+ Specifies the event(s) that this flow depends on.
1449
+
1450
+ ```
1451
+ @trigger(event='foo')
1452
+ ```
1453
+ or
1454
+ ```
1455
+ @trigger(events=['foo', 'bar'])
1456
+ ```
1457
+
1458
+ Additionally, you can specify the parameter mappings
1459
+ to map event payload to Metaflow parameters for the flow.
1460
+ ```
1461
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1462
+ ```
1463
+ or
1464
+ ```
1465
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1466
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1467
+ ```
1468
+
1469
+ 'parameters' can also be a list of strings and tuples like so:
1470
+ ```
1471
+ @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1472
+ ```
1473
+ This is equivalent to:
1474
+ ```
1475
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1476
+ ```
1477
+
1478
+
1479
+ Parameters
1480
+ ----------
1481
+ event : Union[str, Dict[str, Any]], optional, default None
1482
+ Event dependency for this flow.
1483
+ events : List[Union[str, Dict[str, Any]]], default []
1484
+ Events dependency for this flow.
1485
+ options : Dict[str, Any], default {}
1486
+ Backend-specific configuration for tuning eventing behavior.
1487
+ """
1488
+ ...
1489
+
1490
1490
  @typing.overload
1491
1491
  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]]:
1492
1492
  """