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