ob-metaflow-stubs 6.0.3.153__py2.py3-none-any.whl → 6.0.3.155rc0__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 +713 -598
  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 +4 -4
  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 +5 -5
  12. metaflow-stubs/generated_for.txt +1 -1
  13. metaflow-stubs/includefile.pyi +4 -4
  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 +98 -98
  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 +4 -4
  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 +7 -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 +11 -5
  36. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/decorator.pyi +6 -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 +4 -2
  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 +95 -0
  42. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/core.pyi +4 -8
  43. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/decorator.pyi +135 -0
  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 +33 -5
  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 +3 -3
  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 +4 -4
  78. metaflow-stubs/plugins/__init__.pyi +15 -15
  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 +9 -3
  90. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +3 -3
  91. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +5 -5
  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 +4 -4
  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 +5 -5
  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 +4 -4
  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 +6 -6
  117. metaflow-stubs/plugins/cards/card_client.pyi +2 -2
  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 +5 -5
  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 +4 -4
  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 +2 -2
  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 +3 -3
  162. metaflow-stubs/plugins/pypi/conda_decorator.pyi +2 -2
  163. metaflow-stubs/plugins/pypi/conda_environment.pyi +6 -6
  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 +4 -4
  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 +5 -5
  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 +4 -4
  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 +5 -5
  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.153.dist-info → ob_metaflow_stubs-6.0.3.155rc0.dist-info}/METADATA +1 -1
  197. ob_metaflow_stubs-6.0.3.155rc0.dist-info/RECORD +200 -0
  198. ob_metaflow_stubs-6.0.3.153.dist-info/RECORD +0 -198
  199. {ob_metaflow_stubs-6.0.3.153.dist-info → ob_metaflow_stubs-6.0.3.155rc0.dist-info}/WHEEL +0 -0
  200. {ob_metaflow_stubs-6.0.3.153.dist-info → ob_metaflow_stubs-6.0.3.155rc0.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.1+obcheckpoint(0.1.9);ob(v1) #
4
- # Generated on 2025-03-25T18:39:51.892760 #
3
+ # MF version: 2.15.7.2+obcheckpoint(0.2.0);ob(v1) #
4
+ # Generated on 2025-04-14T19:55:43.556579 #
5
5
  ######################################################################################################
6
6
 
7
7
  from __future__ import annotations
8
8
 
9
9
  import typing
10
10
  if typing.TYPE_CHECKING:
11
- import typing
12
11
  import datetime
12
+ import typing
13
13
  FlowSpecDerived = typing.TypeVar("FlowSpecDerived", bound="FlowSpec", contravariant=False, covariant=False)
14
14
  StepFlag = typing.NewType("StepFlag", bool)
15
15
 
@@ -35,16 +35,16 @@ from .user_configs.config_parameters import ConfigValue as ConfigValue
35
35
  from .user_configs.config_parameters import config_expr as config_expr
36
36
  from .user_configs.config_decorators import CustomFlowDecorator as CustomFlowDecorator
37
37
  from .user_configs.config_decorators import CustomStepDecorator as CustomStepDecorator
38
+ from . import events as events
38
39
  from . import tuple_util as tuple_util
39
40
  from . import cards as cards
40
- from . import events as events
41
41
  from . import runner as runner
42
42
  from . import plugins as plugins
43
43
  from .mf_extensions.outerbounds.toplevel.global_aliases_for_metaflow_package import S3 as S3
44
44
  from . import includefile as includefile
45
45
  from .includefile import IncludeFile as IncludeFile
46
- from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
47
46
  from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
47
+ from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_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
@@ -66,6 +66,7 @@ from .runner.deployer import DeployedFlow as DeployedFlow
66
66
  from .runner.nbdeploy import NBDeployer as NBDeployer
67
67
  from .mf_extensions.obcheckpoint.plugins.machine_learning_utilities.checkpoints.final_api import Checkpoint as Checkpoint
68
68
  from .mf_extensions.obcheckpoint.plugins.machine_learning_utilities.datastructures import load_model as load_model
69
+ from .mf_extensions.obcheckpoint.plugins.machine_learning_utilities.datastore.context import artifact_store_from as artifact_store_from
69
70
  from .mf_extensions.outerbounds.toplevel.global_aliases_for_metaflow_package import get_aws_client as get_aws_client
70
71
  from .mf_extensions.outerbounds.plugins.snowflake.snowflake import Snowflake as Snowflake
71
72
  from . import cli_components as cli_components
@@ -150,133 +151,88 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
150
151
  ...
151
152
 
152
153
  @typing.overload
153
- 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]]]:
154
+ 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]]]:
154
155
  """
155
- Creates a human-readable report, a Metaflow Card, after this step completes.
156
+ Specifies the PyPI packages for the step.
156
157
 
157
- Note that you may add multiple `@card` decorators in a step with different parameters.
158
+ Information in this decorator will augment any
159
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
160
+ you can use `@pypi_base` to set packages required by all
161
+ steps and use `@pypi` to specify step-specific overrides.
158
162
 
159
163
 
160
164
  Parameters
161
165
  ----------
162
- type : str, default 'default'
163
- Card type.
164
- id : str, optional, default None
165
- If multiple cards are present, use this id to identify this card.
166
- options : Dict[str, Any], default {}
167
- Options passed to the card. The contents depend on the card type.
168
- timeout : int, default 45
169
- Interrupt reporting if it takes more than this many seconds.
166
+ packages : Dict[str, str], default: {}
167
+ Packages to use for this step. The key is the name of the package
168
+ and the value is the version to use.
169
+ python : str, optional, default: None
170
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
171
+ that the version used will correspond to the version of the Python interpreter used to start the run.
170
172
  """
171
173
  ...
172
174
 
173
175
  @typing.overload
174
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
176
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
175
177
  ...
176
178
 
177
179
  @typing.overload
178
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
180
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
179
181
  ...
180
182
 
181
- 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):
183
+ def pypi(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
182
184
  """
183
- Creates a human-readable report, a Metaflow Card, after this step completes.
185
+ Specifies the PyPI packages for the step.
184
186
 
185
- Note that you may add multiple `@card` decorators in a step with different parameters.
187
+ Information in this decorator will augment any
188
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
189
+ you can use `@pypi_base` to set packages required by all
190
+ steps and use `@pypi` to specify step-specific overrides.
186
191
 
187
192
 
188
193
  Parameters
189
194
  ----------
190
- type : str, default 'default'
191
- Card type.
192
- id : str, optional, default None
193
- If multiple cards are present, use this id to identify this card.
194
- options : Dict[str, Any], default {}
195
- Options passed to the card. The contents depend on the card type.
196
- timeout : int, default 45
197
- Interrupt reporting if it takes more than this many seconds.
195
+ packages : Dict[str, str], default: {}
196
+ Packages to use for this step. The key is the name of the package
197
+ and the value is the version to use.
198
+ python : str, optional, default: None
199
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
200
+ that the version used will correspond to the version of the Python interpreter used to start the run.
198
201
  """
199
202
  ...
200
203
 
201
204
  @typing.overload
202
- 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]]]:
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]]]:
203
206
  """
204
- Enables checkpointing for a step.
205
-
207
+ Specifies secrets to be retrieved and injected as environment variables prior to
208
+ the execution of a step.
206
209
 
207
210
 
208
211
  Parameters
209
212
  ----------
210
- load_policy : str, default: "fresh"
211
- The policy for loading the checkpoint. The following policies are supported:
212
- - "eager": Loads the the latest available checkpoint within the namespace.
213
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
214
- will be loaded at the start of the task.
215
- - "none": Do not load any checkpoint
216
- - "fresh": Loads the lastest checkpoint created within the running Task.
217
- This mode helps loading checkpoints across various retry attempts of the same task.
218
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
219
- created within the task will be loaded when the task is retries execution on failure.
220
-
221
- temp_dir_root : str, default: None
222
- The root directory under which `current.checkpoint.directory` will be created.
213
+ sources : List[Union[str, Dict[str, Any]]], default: []
214
+ List of secret specs, defining how the secrets are to be retrieved
223
215
  """
224
216
  ...
225
217
 
226
218
  @typing.overload
227
- def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
219
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
228
220
  ...
229
221
 
230
222
  @typing.overload
231
- def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
232
- ...
233
-
234
- 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):
235
- """
236
- Enables checkpointing for a step.
237
-
238
-
239
-
240
- Parameters
241
- ----------
242
- load_policy : str, default: "fresh"
243
- The policy for loading the checkpoint. The following policies are supported:
244
- - "eager": Loads the the latest available checkpoint within the namespace.
245
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
246
- will be loaded at the start of the task.
247
- - "none": Do not load any checkpoint
248
- - "fresh": Loads the lastest checkpoint created within the running Task.
249
- This mode helps loading checkpoints across various retry attempts of the same task.
250
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
251
- created within the task will be loaded when the task is retries execution on failure.
252
-
253
- temp_dir_root : str, default: None
254
- The root directory under which `current.checkpoint.directory` will be created.
255
- """
223
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
256
224
  ...
257
225
 
258
- 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]]]:
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]]] = []):
259
227
  """
260
- Decorator that helps cache, version and store models/datasets from huggingface hub.
228
+ Specifies secrets to be retrieved and injected as environment variables prior to
229
+ the execution of a step.
261
230
 
262
231
 
263
232
  Parameters
264
233
  ----------
265
- temp_dir_root : str, optional
266
- The root directory that will hold the temporary directory where objects will be downloaded.
267
-
268
- load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
269
- The list of repos (models/datasets) to load.
270
-
271
- Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
272
-
273
- - If repo (model/dataset) is not found in the datastore:
274
- - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
275
- - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
276
- - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
277
-
278
- - If repo is found in the datastore:
279
- - Loads it directly from datastore to local path (can be temporary directory or specified path)
234
+ sources : List[Union[str, Dict[str, Any]]], default: []
235
+ List of secret specs, defining how the secrets are to be retrieved
280
236
  """
281
237
  ...
282
238
 
@@ -335,160 +291,168 @@ def retry(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
335
291
  """
336
292
  ...
337
293
 
338
- @typing.overload
339
- 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]]]:
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]]]:
340
295
  """
341
- Enables loading / saving of models within a step.
342
-
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.
343
298
 
344
299
 
345
300
  Parameters
346
301
  ----------
347
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
348
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
349
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
350
- - `current.checkpoint`
351
- - `current.model`
352
- - `current.huggingface_hub`
302
+ app_port : int
303
+ Number of GPUs to use.
304
+ app_name : str
305
+ Name of the app to deploy.
306
+ """
307
+ ...
308
+
309
+ @typing.overload
310
+ def timeout(*, seconds: int = 0, minutes: int = 0, hours: int = 0) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
311
+ """
312
+ Specifies a timeout for your step.
353
313
 
354
- 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
355
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
356
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
314
+ This decorator is useful if this step may hang indefinitely.
357
315
 
358
- temp_dir_root : str, default: None
359
- The root directory under which `current.model.loaded` will store loaded models
316
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
317
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
318
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
319
+
320
+ Note that all the values specified in parameters are added together so if you specify
321
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
322
+
323
+
324
+ Parameters
325
+ ----------
326
+ seconds : int, default 0
327
+ Number of seconds to wait prior to timing out.
328
+ minutes : int, default 0
329
+ Number of minutes to wait prior to timing out.
330
+ hours : int, default 0
331
+ Number of hours to wait prior to timing out.
360
332
  """
361
333
  ...
362
334
 
363
335
  @typing.overload
364
- def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
336
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
365
337
  ...
366
338
 
367
339
  @typing.overload
368
- def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
340
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
369
341
  ...
370
342
 
371
- 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):
343
+ def timeout(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, seconds: int = 0, minutes: int = 0, hours: int = 0):
372
344
  """
373
- Enables loading / saving of models within a step.
345
+ Specifies a timeout for your step.
374
346
 
347
+ This decorator is useful if this step may hang indefinitely.
375
348
 
349
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
350
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
351
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
376
352
 
377
- Parameters
378
- ----------
379
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
380
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
381
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
382
- - `current.checkpoint`
383
- - `current.model`
384
- - `current.huggingface_hub`
353
+ Note that all the values specified in parameters are added together so if you specify
354
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
385
355
 
386
- 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
387
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
388
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
389
356
 
390
- temp_dir_root : str, default: None
391
- The root directory under which `current.model.loaded` will store loaded models
357
+ Parameters
358
+ ----------
359
+ seconds : int, default 0
360
+ Number of seconds to wait prior to timing out.
361
+ minutes : int, default 0
362
+ Number of minutes to wait prior to timing out.
363
+ hours : int, default 0
364
+ Number of hours to wait prior to timing out.
392
365
  """
393
366
  ...
394
367
 
395
- 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]]]:
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]]]:
396
369
  """
397
- This decorator is used to run NIM containers in Metaflow tasks as sidecars.
370
+ Decorator that helps cache, version and store models/datasets from huggingface hub.
398
371
 
399
- User code call
400
- -----------
401
- @nim(
402
- models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
403
- backend='managed'
404
- )
405
372
 
406
- Valid backend options
407
- ---------------------
408
- - 'managed': Outerbounds selects a compute provider based on the model.
373
+ Parameters
374
+ ----------
375
+ temp_dir_root : str, optional
376
+ The root directory that will hold the temporary directory where objects will be downloaded.
409
377
 
410
- Valid model options
411
- ----------------
412
- - 'meta/llama3-8b-instruct': 8B parameter model
413
- - 'meta/llama3-70b-instruct': 70B parameter model
414
- - any model here: https://nvcf.ngc.nvidia.com/functions?filter=nvidia-functions
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
+ 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
+ """
395
+ Specifies that this step should execute on DGX cloud.
415
396
 
416
397
 
417
398
  Parameters
418
399
  ----------
419
- models: list[NIM]
420
- List of NIM containers running models in sidecars.
421
- backend: str
422
- Compute provider to run the NIM container.
400
+ gpu : int
401
+ Number of GPUs to use.
402
+ gpu_type : str
403
+ Type of Nvidia GPU to use.
423
404
  queue_timeout : int
424
405
  Time to keep the job in NVCF's queue.
425
406
  """
426
407
  ...
427
408
 
428
409
  @typing.overload
429
- 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]]]:
410
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
430
411
  """
431
- Specifies that the step will success under all circumstances.
432
-
433
- The decorator will create an optional artifact, specified by `var`, which
434
- contains the exception raised. You can use it to detect the presence
435
- of errors, indicating that all happy-path artifacts produced by the step
436
- are missing.
437
-
438
-
439
- Parameters
440
- ----------
441
- var : str, optional, default None
442
- Name of the artifact in which to store the caught exception.
443
- If not specified, the exception is not stored.
444
- print_exception : bool, default True
445
- Determines whether or not the exception is printed to
446
- stdout when caught.
412
+ Internal decorator to support Fast bakery
447
413
  """
448
414
  ...
449
415
 
450
416
  @typing.overload
451
- def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
417
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
452
418
  ...
453
419
 
454
- @typing.overload
455
- def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
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
+ """
456
424
  ...
457
425
 
458
- 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):
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]]]:
459
427
  """
460
- Specifies that the step will success under all circumstances.
428
+ This decorator is used to run Ollama APIs as Metaflow task sidecars.
461
429
 
462
- The decorator will create an optional artifact, specified by `var`, which
463
- contains the exception raised. You can use it to detect the presence
464
- of errors, indicating that all happy-path artifacts produced by the step
465
- are missing.
430
+ User code call
431
+ -----------
432
+ @ollama(
433
+ models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
434
+ backend='local'
435
+ )
466
436
 
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.
467
442
 
468
- Parameters
469
- ----------
470
- var : str, optional, default None
471
- Name of the artifact in which to store the caught exception.
472
- If not specified, the exception is not stored.
473
- print_exception : bool, default True
474
- Determines whether or not the exception is printed to
475
- stdout when caught.
476
- """
477
- ...
478
-
479
- def nvidia(*, gpu: int, gpu_type: str, queue_timeout: int) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
480
- """
481
- Specifies that this step should execute on DGX cloud.
443
+ Valid model options
444
+ ----------------
445
+ - 'llama3.2'
446
+ - 'llama3.3'
447
+ - any model here https://ollama.com/search
482
448
 
483
449
 
484
450
  Parameters
485
451
  ----------
486
- gpu : int
487
- Number of GPUs to use.
488
- gpu_type : str
489
- Type of Nvidia GPU to use.
490
- queue_timeout : int
491
- Time to keep the job in NVCF's queue.
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.
492
456
  """
493
457
  ...
494
458
 
@@ -551,137 +515,93 @@ def conda(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
551
515
  """
552
516
  ...
553
517
 
554
- 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]]]:
518
+ @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
520
  """
556
- Specifies that this step is used to deploy an instance of the app.
557
- Requires that self.app_name, self.app_port, self.entrypoint and self.deployDir is set.
521
+ Enables loading / saving of models within a step.
522
+
558
523
 
559
524
 
560
525
  Parameters
561
526
  ----------
562
- app_port : int
563
- Number of GPUs to use.
564
- app_name : str
565
- Name of the app to deploy.
566
- """
567
- ...
568
-
569
- @typing.overload
570
- 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]]]:
571
- """
572
- Specifies secrets to be retrieved and injected as environment variables prior to
573
- the execution of a step.
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`
574
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.
575
537
 
576
- Parameters
577
- ----------
578
- sources : List[Union[str, Dict[str, Any]]], default: []
579
- List of secret specs, defining how the secrets are to be retrieved
538
+ temp_dir_root : str, default: None
539
+ The root directory under which `current.model.loaded` will store loaded models
580
540
  """
581
541
  ...
582
542
 
583
543
  @typing.overload
584
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
544
+ def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
585
545
  ...
586
546
 
587
547
  @typing.overload
588
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
548
+ def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
589
549
  ...
590
550
 
591
- 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]]] = []):
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):
592
552
  """
593
- Specifies secrets to be retrieved and injected as environment variables prior to
594
- the execution of a step.
553
+ Enables loading / saving of models within a step.
554
+
595
555
 
596
556
 
597
557
  Parameters
598
558
  ----------
599
- sources : List[Union[str, Dict[str, Any]]], default: []
600
- List of secret specs, defining how the secrets are to be retrieved
601
- """
602
- ...
603
-
604
- 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]]]:
605
- """
606
- This decorator is used to run Ollama APIs as Metaflow task sidecars.
607
-
608
- User code call
609
- -----------
610
- @ollama(
611
- models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
612
- backend='local'
613
- )
614
-
615
- Valid backend options
616
- ---------------------
617
- - 'local': Run as a separate process on the local task machine.
618
- - (TODO) 'managed': Outerbounds hosts and selects compute provider.
619
- - (TODO) 'remote': Spin up separate instance to serve Ollama models.
620
-
621
- Valid model options
622
- ----------------
623
- - 'llama3.2'
624
- - 'llama3.3'
625
- - any model here https://ollama.com/search
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`
626
565
 
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.
627
569
 
628
- Parameters
629
- ----------
630
- models: list[Ollama]
631
- List of Ollama containers running models in sidecars.
632
- backend: str
633
- Determines where and how to run the Ollama process.
570
+ temp_dir_root : str, default: None
571
+ The root directory under which `current.model.loaded` will store loaded models
634
572
  """
635
573
  ...
636
574
 
637
575
  @typing.overload
638
- 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]]]:
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]]]:
639
577
  """
640
- Specifies the PyPI packages for the step.
641
-
642
- Information in this decorator will augment any
643
- attributes set in the `@pyi_base` flow-level decorator. Hence,
644
- you can use `@pypi_base` to set packages required by all
645
- steps and use `@pypi` to specify step-specific overrides.
578
+ Specifies environment variables to be set prior to the execution of a step.
646
579
 
647
580
 
648
581
  Parameters
649
582
  ----------
650
- packages : Dict[str, str], default: {}
651
- Packages to use for this step. The key is the name of the package
652
- and the value is the version to use.
653
- python : str, optional, default: None
654
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
655
- that the version used will correspond to the version of the Python interpreter used to start the run.
583
+ vars : Dict[str, str], default {}
584
+ Dictionary of environment variables to set.
656
585
  """
657
586
  ...
658
587
 
659
588
  @typing.overload
660
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
589
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
661
590
  ...
662
591
 
663
592
  @typing.overload
664
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
593
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
665
594
  ...
666
595
 
667
- def pypi(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
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] = {}):
668
597
  """
669
- Specifies the PyPI packages for the step.
670
-
671
- Information in this decorator will augment any
672
- attributes set in the `@pyi_base` flow-level decorator. Hence,
673
- you can use `@pypi_base` to set packages required by all
674
- steps and use `@pypi` to specify step-specific overrides.
598
+ Specifies environment variables to be set prior to the execution of a step.
675
599
 
676
600
 
677
601
  Parameters
678
602
  ----------
679
- packages : Dict[str, str], default: {}
680
- Packages to use for this step. The key is the name of the package
681
- and the value is the version to use.
682
- python : str, optional, default: None
683
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
684
- that the version used will correspond to the version of the Python interpreter used to start the run.
603
+ vars : Dict[str, str], default {}
604
+ Dictionary of environment variables to set.
685
605
  """
686
606
  ...
687
607
 
@@ -764,79 +684,55 @@ def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None]
764
684
  """
765
685
  ...
766
686
 
767
- @typing.overload
768
- def timeout(*, seconds: int = 0, minutes: int = 0, hours: int = 0) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
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]]]:
769
688
  """
770
- Specifies a timeout for your step.
689
+ This decorator is used to run NIM containers in Metaflow tasks as sidecars.
771
690
 
772
- This decorator is useful if this step may hang indefinitely.
691
+ User code call
692
+ -----------
693
+ @nim(
694
+ models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
695
+ backend='managed'
696
+ )
773
697
 
774
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
775
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
776
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
698
+ Valid backend options
699
+ ---------------------
700
+ - 'managed': Outerbounds selects a compute provider based on the model.
777
701
 
778
- Note that all the values specified in parameters are added together so if you specify
779
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
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
780
707
 
781
708
 
782
709
  Parameters
783
710
  ----------
784
- seconds : int, default 0
785
- Number of seconds to wait prior to timing out.
786
- minutes : int, default 0
787
- Number of minutes to wait prior to timing out.
788
- hours : int, default 0
789
- Number of hours to wait prior to timing out.
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.
790
717
  """
791
718
  ...
792
719
 
793
720
  @typing.overload
794
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
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
+ """
795
726
  ...
796
727
 
797
728
  @typing.overload
798
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
729
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
799
730
  ...
800
731
 
801
- def timeout(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, seconds: int = 0, minutes: int = 0, hours: int = 0):
732
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
802
733
  """
803
- Specifies a timeout for your step.
804
-
805
- This decorator is useful if this step may hang indefinitely.
806
-
807
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
808
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
809
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
810
-
811
- Note that all the values specified in parameters are added together so if you specify
812
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
813
-
814
-
815
- Parameters
816
- ----------
817
- seconds : int, default 0
818
- Number of seconds to wait prior to timing out.
819
- minutes : int, default 0
820
- Number of minutes to wait prior to timing out.
821
- hours : int, default 0
822
- Number of hours to wait prior to timing out.
823
- """
824
- ...
825
-
826
- @typing.overload
827
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
828
- """
829
- Internal decorator to support Fast bakery
830
- """
831
- ...
832
-
833
- @typing.overload
834
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
835
- ...
836
-
837
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
838
- """
839
- Internal decorator to support Fast bakery
734
+ Decorator prototype for all step decorators. This function gets specialized
735
+ and imported for all decorators types by _import_plugin_decorators().
840
736
  """
841
737
  ...
842
738
 
@@ -918,54 +814,159 @@ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: ty
918
814
  ...
919
815
 
920
816
  @typing.overload
921
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
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]]]:
922
818
  """
923
- Decorator prototype for all step decorators. This function gets specialized
924
- and imported for all decorators types by _import_plugin_decorators().
819
+ Creates a human-readable report, a Metaflow Card, after this step completes.
820
+
821
+ Note that you may add multiple `@card` decorators in a step with different parameters.
822
+
823
+
824
+ Parameters
825
+ ----------
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.
925
834
  """
926
835
  ...
927
836
 
928
837
  @typing.overload
929
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
838
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
930
839
  ...
931
840
 
932
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
841
+ @typing.overload
842
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
843
+ ...
844
+
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):
933
846
  """
934
- Decorator prototype for all step decorators. This function gets specialized
935
- and imported for all decorators types by _import_plugin_decorators().
847
+ Creates a human-readable report, a Metaflow Card, after this step completes.
848
+
849
+ Note that you may add multiple `@card` decorators in a step with different parameters.
850
+
851
+
852
+ Parameters
853
+ ----------
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.
936
862
  """
937
863
  ...
938
864
 
939
865
  @typing.overload
940
- 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]]]:
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]]]:
941
867
  """
942
- Specifies environment variables to be set prior to the execution of a step.
868
+ Enables checkpointing for a step.
869
+
943
870
 
944
871
 
945
872
  Parameters
946
873
  ----------
947
- vars : Dict[str, str], default {}
948
- Dictionary of environment variables to set.
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.
949
887
  """
950
888
  ...
951
889
 
952
890
  @typing.overload
953
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
891
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
954
892
  ...
955
893
 
956
894
  @typing.overload
957
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
895
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
958
896
  ...
959
897
 
960
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
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):
961
899
  """
962
- Specifies environment variables to be set prior to the execution of a step.
900
+ Enables checkpointing for a step.
901
+
963
902
 
964
903
 
965
904
  Parameters
966
905
  ----------
967
- vars : Dict[str, str], default {}
968
- Dictionary of environment variables to set.
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.
919
+ """
920
+ ...
921
+
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]]]:
924
+ """
925
+ Specifies that the step will success under all circumstances.
926
+
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.
931
+
932
+
933
+ Parameters
934
+ ----------
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.
941
+ """
942
+ ...
943
+
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):
953
+ """
954
+ Specifies that the step will success under all circumstances.
955
+
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.
960
+
961
+
962
+ Parameters
963
+ ----------
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.
969
970
  """
970
971
  ...
971
972
 
@@ -1020,97 +1021,311 @@ def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packa
1020
1021
  """
1021
1022
  ...
1022
1023
 
1023
- 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]]:
1024
+ @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]]:
1024
1026
  """
1025
- 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)
1026
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1027
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1028
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
1029
- starts only after all sensors finish.
1027
+ Specifies the event(s) that this flow depends on.
1028
+
1029
+ ```
1030
+ @trigger(event='foo')
1031
+ ```
1032
+ or
1033
+ ```
1034
+ @trigger(events=['foo', 'bar'])
1035
+ ```
1036
+
1037
+ Additionally, you can specify the parameter mappings
1038
+ to map event payload to Metaflow parameters for the flow.
1039
+ ```
1040
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1041
+ ```
1042
+ or
1043
+ ```
1044
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1045
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1046
+ ```
1047
+
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:
1053
+ ```
1054
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1055
+ ```
1030
1056
 
1031
1057
 
1032
1058
  Parameters
1033
1059
  ----------
1034
- timeout : int
1035
- Time, in seconds before the task times out and fails. (Default: 3600)
1036
- poke_interval : int
1037
- Time in seconds that the job should wait in between each try. (Default: 60)
1038
- mode : str
1039
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1040
- exponential_backoff : bool
1041
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1042
- pool : str
1043
- the slot pool this task should run in,
1044
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1045
- soft_fail : bool
1046
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1047
- name : str
1048
- Name of the sensor on Airflow
1049
- description : str
1050
- Description of sensor in the Airflow UI
1051
- bucket_key : Union[str, List[str]]
1052
- The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1053
- When it's specified as a full s3:// url, please leave `bucket_name` as None
1054
- bucket_name : str
1055
- Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1056
- When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1057
- wildcard_match : bool
1058
- whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1059
- aws_conn_id : str
1060
- a reference to the s3 connection on Airflow. (Default: None)
1061
- verify : bool
1062
- Whether or not to verify SSL certificates for S3 connection. (Default: None)
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.
1064
+ options : Dict[str, Any], default {}
1065
+ Backend-specific configuration for tuning eventing behavior.
1063
1066
  """
1064
1067
  ...
1065
1068
 
1066
1069
  @typing.overload
1067
- 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]]:
1070
+ def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1071
+ ...
1072
+
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] = {}):
1068
1074
  """
1069
- Specifies the times when the flow should be run when running on a
1070
- production scheduler.
1071
-
1075
+ Specifies the event(s) that this flow depends on.
1072
1076
 
1073
- Parameters
1074
- ----------
1075
- hourly : bool, default False
1076
- Run the workflow hourly.
1077
- daily : bool, default True
1078
- Run the workflow daily.
1079
- weekly : bool, default False
1080
- Run the workflow weekly.
1081
- cron : str, optional, default None
1082
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1083
- specified by this expression.
1084
- timezone : str, optional, default None
1085
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1086
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1077
+ ```
1078
+ @trigger(event='foo')
1079
+ ```
1080
+ or
1081
+ ```
1082
+ @trigger(events=['foo', 'bar'])
1083
+ ```
1084
+
1085
+ Additionally, you can specify the parameter mappings
1086
+ to map event payload to Metaflow parameters for the flow.
1087
+ ```
1088
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1089
+ ```
1090
+ or
1091
+ ```
1092
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1093
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1094
+ ```
1095
+
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:
1101
+ ```
1102
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1103
+ ```
1104
+
1105
+
1106
+ Parameters
1107
+ ----------
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.
1112
+ options : Dict[str, Any], default {}
1113
+ Backend-specific configuration for tuning eventing behavior.
1087
1114
  """
1088
1115
  ...
1089
1116
 
1090
1117
  @typing.overload
1091
- def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
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]]:
1119
+ """
1120
+ Specifies the flow(s) that this flow depends on.
1121
+
1122
+ ```
1123
+ @trigger_on_finish(flow='FooFlow')
1124
+ ```
1125
+ or
1126
+ ```
1127
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1128
+ ```
1129
+ This decorator respects the @project decorator and triggers the flow
1130
+ when upstream runs within the same namespace complete successfully
1131
+
1132
+ Additionally, you can specify project aware upstream flow dependencies
1133
+ by specifying the fully qualified project_flow_name.
1134
+ ```
1135
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1136
+ ```
1137
+ or
1138
+ ```
1139
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1140
+ ```
1141
+
1142
+ You can also specify just the project or project branch (other values will be
1143
+ inferred from the current project or project branch):
1144
+ ```
1145
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1146
+ ```
1147
+
1148
+ Note that `branch` is typically one of:
1149
+ - `prod`
1150
+ - `user.bob`
1151
+ - `test.my_experiment`
1152
+ - `prod.staging`
1153
+
1154
+
1155
+ Parameters
1156
+ ----------
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.
1161
+ options : Dict[str, Any], default {}
1162
+ Backend-specific configuration for tuning eventing behavior.
1163
+ """
1092
1164
  ...
1093
1165
 
1094
- 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):
1166
+ @typing.overload
1167
+ def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1168
+ ...
1169
+
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] = {}):
1095
1171
  """
1096
- Specifies the times when the flow should be run when running on a
1097
- production scheduler.
1172
+ Specifies the flow(s) that this flow depends on.
1173
+
1174
+ ```
1175
+ @trigger_on_finish(flow='FooFlow')
1176
+ ```
1177
+ or
1178
+ ```
1179
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1180
+ ```
1181
+ This decorator respects the @project decorator and triggers the flow
1182
+ when upstream runs within the same namespace complete successfully
1183
+
1184
+ Additionally, you can specify project aware upstream flow dependencies
1185
+ by specifying the fully qualified project_flow_name.
1186
+ ```
1187
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1188
+ ```
1189
+ or
1190
+ ```
1191
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1192
+ ```
1193
+
1194
+ You can also specify just the project or project branch (other values will be
1195
+ inferred from the current project or project branch):
1196
+ ```
1197
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1198
+ ```
1199
+
1200
+ Note that `branch` is typically one of:
1201
+ - `prod`
1202
+ - `user.bob`
1203
+ - `test.my_experiment`
1204
+ - `prod.staging`
1098
1205
 
1099
1206
 
1100
1207
  Parameters
1101
1208
  ----------
1102
- hourly : bool, default False
1103
- Run the workflow hourly.
1104
- daily : bool, default True
1105
- Run the workflow daily.
1106
- weekly : bool, default False
1107
- Run the workflow weekly.
1108
- cron : str, optional, default None
1109
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1110
- specified by this expression.
1111
- timezone : str, optional, default None
1112
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1113
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
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.
1215
+ """
1216
+ ...
1217
+
1218
+ def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None):
1219
+ """
1220
+ Allows setting external datastores to save data for the
1221
+ `@checkpoint`/`@model`/`@huggingface_hub` decorators.
1222
+
1223
+ This decorator is useful when users wish to save data to a different datastore
1224
+ than what is configured in Metaflow. This can be for variety of reasons:
1225
+
1226
+ 1. Data security: The objects needs to be stored in a bucket (object storage) that is not accessible by other flows.
1227
+ 2. Data Locality: The location where the task is executing is not located in the same region as the datastore.
1228
+ - Example: Metaflow datastore lives in US East, but the task is executing in Finland datacenters.
1229
+ 3. Data Lifecycle Policies: The objects need to be archived / managed separately from the Metaflow managed objects.
1230
+ - Example: Flow is training very large models that need to be stored separately and will be deleted more aggressively than the Metaflow managed objects.
1231
+
1232
+ Usage:
1233
+ ----------
1234
+
1235
+ - Using a custom IAM role to access the datastore.
1236
+
1237
+ ```python
1238
+ @with_artifact_store(
1239
+ type="s3",
1240
+ config=lambda: {
1241
+ "root": "s3://my-bucket-foo/path/to/root",
1242
+ "role_arn": ROLE,
1243
+ },
1244
+ )
1245
+ class MyFlow(FlowSpec):
1246
+
1247
+ @checkpoint
1248
+ @step
1249
+ def start(self):
1250
+ with open("my_file.txt", "w") as f:
1251
+ f.write("Hello, World!")
1252
+ self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1253
+ self.next(self.end)
1254
+
1255
+ ```
1256
+
1257
+ - Using credentials to access the s3-compatible datastore.
1258
+
1259
+ ```python
1260
+ @with_artifact_store(
1261
+ type="s3",
1262
+ config=lambda: {
1263
+ "root": "s3://my-bucket-foo/path/to/root",
1264
+ "client_params": {
1265
+ "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1266
+ "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1267
+ },
1268
+ },
1269
+ )
1270
+ class MyFlow(FlowSpec):
1271
+
1272
+ @checkpoint
1273
+ @step
1274
+ def start(self):
1275
+ with open("my_file.txt", "w") as f:
1276
+ f.write("Hello, World!")
1277
+ self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1278
+ self.next(self.end)
1279
+
1280
+ ```
1281
+
1282
+ - Accessing objects stored in external datastores after task execution.
1283
+
1284
+ ```python
1285
+ run = Run("CheckpointsTestsFlow/8992")
1286
+ with artifact_store_from(run=run, config={
1287
+ "client_params": {
1288
+ "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1289
+ "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1290
+ },
1291
+ }):
1292
+ with Checkpoint() as cp:
1293
+ latest = cp.list(
1294
+ task=run["start"].task
1295
+ )[0]
1296
+ print(latest)
1297
+ cp.load(
1298
+ latest,
1299
+ "test-checkpoints"
1300
+ )
1301
+
1302
+ task = Task("TorchTuneFlow/8484/train/53673")
1303
+ with artifact_store_from(run=run, config={
1304
+ "client_params": {
1305
+ "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1306
+ "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1307
+ },
1308
+ }):
1309
+ load_model(
1310
+ task.data.model_ref,
1311
+ "test-models"
1312
+ )
1313
+ ```
1314
+ Parameters:
1315
+ ----------
1316
+
1317
+ type: str
1318
+ The type of the datastore. Can be one of 's3', 'gcs', 'azure' or any other supported metaflow Datastore.
1319
+
1320
+ config: dict or Callable
1321
+ Dictionary of configuration options for the datastore. The following keys are required:
1322
+ - root: The root path in the datastore where the data will be saved. (needs to be in the format expected by the datastore)
1323
+ - example: 's3://bucket-name/path/to/root'
1324
+ - example: 'gs://bucket-name/path/to/root'
1325
+ - example: 'https://myblockacc.blob.core.windows.net/metaflow/'
1326
+ - role_arn (optional): AWS IAM role to access s3 bucket (only when `type` is 's3')
1327
+ - session_vars (optional): AWS session variables to access s3 bucket (only when `type` is 's3')
1328
+ - client_params (optional): AWS client parameters to access s3 bucket (only when `type` is 's3')
1114
1329
  """
1115
1330
  ...
1116
1331
 
@@ -1149,6 +1364,57 @@ def project(*, name: str, branch: typing.Optional[str] = None, production: bool
1149
1364
  """
1150
1365
  ...
1151
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
+ @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):
1396
+ """
1397
+ Specifies the times when the flow should be run when running on a
1398
+ production scheduler.
1399
+
1400
+
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
+
1152
1418
  @typing.overload
1153
1419
  def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1154
1420
  """
@@ -1233,197 +1499,46 @@ def airflow_external_task_sensor(*, timeout: int, poke_interval: int, mode: str,
1233
1499
  """
1234
1500
  ...
1235
1501
 
1236
- @typing.overload
1237
- 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]]:
1238
- """
1239
- Specifies the flow(s) that this flow depends on.
1240
-
1241
- ```
1242
- @trigger_on_finish(flow='FooFlow')
1243
- ```
1244
- or
1245
- ```
1246
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1247
- ```
1248
- This decorator respects the @project decorator and triggers the flow
1249
- when upstream runs within the same namespace complete successfully
1250
-
1251
- Additionally, you can specify project aware upstream flow dependencies
1252
- by specifying the fully qualified project_flow_name.
1253
- ```
1254
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1255
- ```
1256
- or
1257
- ```
1258
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1259
- ```
1260
-
1261
- You can also specify just the project or project branch (other values will be
1262
- inferred from the current project or project branch):
1263
- ```
1264
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1265
- ```
1266
-
1267
- Note that `branch` is typically one of:
1268
- - `prod`
1269
- - `user.bob`
1270
- - `test.my_experiment`
1271
- - `prod.staging`
1272
-
1273
-
1274
- Parameters
1275
- ----------
1276
- flow : Union[str, Dict[str, str]], optional, default None
1277
- Upstream flow dependency for this flow.
1278
- flows : List[Union[str, Dict[str, str]]], default []
1279
- Upstream flow dependencies for this flow.
1280
- options : Dict[str, Any], default {}
1281
- Backend-specific configuration for tuning eventing behavior.
1282
- """
1283
- ...
1284
-
1285
- @typing.overload
1286
- def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1287
- ...
1288
-
1289
- 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] = {}):
1290
- """
1291
- Specifies the flow(s) that this flow depends on.
1292
-
1293
- ```
1294
- @trigger_on_finish(flow='FooFlow')
1295
- ```
1296
- or
1297
- ```
1298
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1299
- ```
1300
- This decorator respects the @project decorator and triggers the flow
1301
- when upstream runs within the same namespace complete successfully
1302
-
1303
- Additionally, you can specify project aware upstream flow dependencies
1304
- by specifying the fully qualified project_flow_name.
1305
- ```
1306
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1307
- ```
1308
- or
1309
- ```
1310
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1311
- ```
1312
-
1313
- You can also specify just the project or project branch (other values will be
1314
- inferred from the current project or project branch):
1315
- ```
1316
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1317
- ```
1318
-
1319
- Note that `branch` is typically one of:
1320
- - `prod`
1321
- - `user.bob`
1322
- - `test.my_experiment`
1323
- - `prod.staging`
1324
-
1325
-
1326
- Parameters
1327
- ----------
1328
- flow : Union[str, Dict[str, str]], optional, default None
1329
- Upstream flow dependency for this flow.
1330
- flows : List[Union[str, Dict[str, str]]], default []
1331
- Upstream flow dependencies for this flow.
1332
- options : Dict[str, Any], default {}
1333
- Backend-specific configuration for tuning eventing behavior.
1334
- """
1335
- ...
1336
-
1337
- @typing.overload
1338
- 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]]:
1339
- """
1340
- Specifies the event(s) that this flow depends on.
1341
-
1342
- ```
1343
- @trigger(event='foo')
1344
- ```
1345
- or
1346
- ```
1347
- @trigger(events=['foo', 'bar'])
1348
- ```
1349
-
1350
- Additionally, you can specify the parameter mappings
1351
- to map event payload to Metaflow parameters for the flow.
1352
- ```
1353
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1354
- ```
1355
- or
1356
- ```
1357
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1358
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1359
- ```
1360
-
1361
- 'parameters' can also be a list of strings and tuples like so:
1362
- ```
1363
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1364
- ```
1365
- This is equivalent to:
1366
- ```
1367
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1368
- ```
1369
-
1370
-
1371
- Parameters
1372
- ----------
1373
- event : Union[str, Dict[str, Any]], optional, default None
1374
- Event dependency for this flow.
1375
- events : List[Union[str, Dict[str, Any]]], default []
1376
- Events dependency for this flow.
1377
- options : Dict[str, Any], default {}
1378
- Backend-specific configuration for tuning eventing behavior.
1379
- """
1380
- ...
1381
-
1382
- @typing.overload
1383
- def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1384
- ...
1385
-
1386
- 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] = {}):
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]]:
1387
1503
  """
1388
- Specifies the event(s) that this flow depends on.
1389
-
1390
- ```
1391
- @trigger(event='foo')
1392
- ```
1393
- or
1394
- ```
1395
- @trigger(events=['foo', 'bar'])
1396
- ```
1397
-
1398
- Additionally, you can specify the parameter mappings
1399
- to map event payload to Metaflow parameters for the flow.
1400
- ```
1401
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1402
- ```
1403
- or
1404
- ```
1405
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1406
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1407
- ```
1408
-
1409
- 'parameters' can also be a list of strings and tuples like so:
1410
- ```
1411
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1412
- ```
1413
- This is equivalent to:
1414
- ```
1415
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1416
- ```
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.
1417
1509
 
1418
1510
 
1419
1511
  Parameters
1420
1512
  ----------
1421
- event : Union[str, Dict[str, Any]], optional, default None
1422
- Event dependency for this flow.
1423
- events : List[Union[str, Dict[str, Any]]], default []
1424
- Events dependency for this flow.
1425
- options : Dict[str, Any], default {}
1426
- Backend-specific configuration for tuning eventing behavior.
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)
1427
1542
  """
1428
1543
  ...
1429
1544