ob-metaflow-stubs 6.0.3.157__py2.py3-none-any.whl → 6.0.3.159__py2.py3-none-any.whl

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