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