ob-metaflow-stubs 6.0.3.182rc2__py2.py3-none-any.whl → 6.0.3.183rc1__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 (220) hide show
  1. metaflow-stubs/__init__.pyi +731 -731
  2. metaflow-stubs/cards.pyi +2 -2
  3. metaflow-stubs/cli.pyi +2 -2
  4. metaflow-stubs/cli_components/__init__.pyi +2 -2
  5. metaflow-stubs/cli_components/utils.pyi +2 -2
  6. metaflow-stubs/client/__init__.pyi +2 -2
  7. metaflow-stubs/client/core.pyi +5 -5
  8. metaflow-stubs/client/filecache.pyi +2 -2
  9. metaflow-stubs/events.pyi +2 -2
  10. metaflow-stubs/exception.pyi +2 -2
  11. metaflow-stubs/flowspec.pyi +5 -5
  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 -4
  20. metaflow-stubs/metaflow_current.pyi +101 -101
  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 +3 -3
  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 +3 -3
  37. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/decorator.pyi +4 -4
  38. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/exceptions.pyi +2 -2
  39. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/final_api.pyi +3 -3
  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 +2 -2
  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 +3 -3
  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 +3 -3
  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/aws/__init__.pyi +2 -2
  67. metaflow-stubs/mf_extensions/outerbounds/plugins/aws/assume_role_decorator.pyi +2 -2
  68. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/__init__.pyi +2 -2
  69. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/async_cards.pyi +2 -2
  70. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/injector.pyi +2 -2
  71. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/__init__.pyi +2 -2
  72. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/coreweave.pyi +2 -2
  73. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/nebius.pyi +2 -2
  74. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/__init__.pyi +2 -2
  75. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/baker.pyi +2 -2
  76. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/docker_environment.pyi +2 -2
  77. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/fast_bakery.pyi +2 -2
  78. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/__init__.pyi +2 -2
  79. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/pod_killer.pyi +2 -2
  80. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/__init__.pyi +2 -2
  81. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/constants.pyi +2 -2
  82. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/exceptions.pyi +2 -2
  83. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/ollama.pyi +2 -2
  84. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/status_card.pyi +2 -2
  85. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/__init__.pyi +2 -2
  86. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/snowflake.pyi +2 -2
  87. metaflow-stubs/mf_extensions/outerbounds/profilers/__init__.pyi +2 -2
  88. metaflow-stubs/mf_extensions/outerbounds/profilers/gpu.pyi +2 -2
  89. metaflow-stubs/mf_extensions/outerbounds/remote_config.pyi +2 -2
  90. metaflow-stubs/mf_extensions/outerbounds/toplevel/__init__.pyi +2 -2
  91. metaflow-stubs/mf_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.pyi +2 -2
  92. metaflow-stubs/multicore_utils.pyi +2 -2
  93. metaflow-stubs/ob_internal.pyi +2 -2
  94. metaflow-stubs/parameters.pyi +3 -3
  95. metaflow-stubs/plugins/__init__.pyi +11 -11
  96. metaflow-stubs/plugins/airflow/__init__.pyi +2 -2
  97. metaflow-stubs/plugins/airflow/airflow_utils.pyi +2 -2
  98. metaflow-stubs/plugins/airflow/exception.pyi +2 -2
  99. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +2 -2
  100. metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi +2 -2
  101. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +2 -2
  102. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +2 -2
  103. metaflow-stubs/plugins/argo/__init__.pyi +2 -2
  104. metaflow-stubs/plugins/argo/argo_client.pyi +2 -2
  105. metaflow-stubs/plugins/argo/argo_events.pyi +2 -2
  106. metaflow-stubs/plugins/argo/argo_workflows.pyi +2 -2
  107. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +3 -3
  108. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +2 -2
  109. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +2 -2
  110. metaflow-stubs/plugins/aws/__init__.pyi +2 -2
  111. metaflow-stubs/plugins/aws/aws_client.pyi +2 -2
  112. metaflow-stubs/plugins/aws/aws_utils.pyi +2 -2
  113. metaflow-stubs/plugins/aws/batch/__init__.pyi +2 -2
  114. metaflow-stubs/plugins/aws/batch/batch.pyi +2 -2
  115. metaflow-stubs/plugins/aws/batch/batch_client.pyi +2 -2
  116. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +2 -2
  117. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +2 -2
  118. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +2 -2
  119. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +2 -2
  120. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +2 -2
  121. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +2 -2
  122. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +2 -2
  123. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +2 -2
  124. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer.pyi +3 -3
  125. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer_objects.pyi +2 -2
  126. metaflow-stubs/plugins/azure/__init__.pyi +2 -2
  127. metaflow-stubs/plugins/azure/azure_credential.pyi +2 -2
  128. metaflow-stubs/plugins/azure/azure_exceptions.pyi +2 -2
  129. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +2 -2
  130. metaflow-stubs/plugins/azure/azure_utils.pyi +2 -2
  131. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +2 -2
  132. metaflow-stubs/plugins/azure/includefile_support.pyi +2 -2
  133. metaflow-stubs/plugins/cards/__init__.pyi +2 -2
  134. metaflow-stubs/plugins/cards/card_client.pyi +3 -3
  135. metaflow-stubs/plugins/cards/card_creator.pyi +2 -2
  136. metaflow-stubs/plugins/cards/card_datastore.pyi +2 -4
  137. metaflow-stubs/plugins/cards/card_decorator.pyi +2 -2
  138. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +2 -2
  139. metaflow-stubs/plugins/cards/card_modules/basic.pyi +3 -3
  140. metaflow-stubs/plugins/cards/card_modules/card.pyi +2 -2
  141. metaflow-stubs/plugins/cards/card_modules/components.pyi +2 -2
  142. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +2 -2
  143. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +2 -2
  144. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +2 -2
  145. metaflow-stubs/plugins/cards/card_resolver.pyi +2 -2
  146. metaflow-stubs/plugins/cards/component_serializer.pyi +2 -2
  147. metaflow-stubs/plugins/cards/exception.pyi +2 -2
  148. metaflow-stubs/plugins/catch_decorator.pyi +3 -3
  149. metaflow-stubs/plugins/datatools/__init__.pyi +2 -2
  150. metaflow-stubs/plugins/datatools/local.pyi +2 -2
  151. metaflow-stubs/plugins/datatools/s3/__init__.pyi +2 -2
  152. metaflow-stubs/plugins/datatools/s3/s3.pyi +2 -2
  153. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +2 -2
  154. metaflow-stubs/plugins/datatools/s3/s3util.pyi +2 -2
  155. metaflow-stubs/plugins/debug_logger.pyi +2 -2
  156. metaflow-stubs/plugins/debug_monitor.pyi +2 -2
  157. metaflow-stubs/plugins/environment_decorator.pyi +2 -2
  158. metaflow-stubs/plugins/events_decorator.pyi +2 -2
  159. metaflow-stubs/plugins/frameworks/__init__.pyi +2 -2
  160. metaflow-stubs/plugins/frameworks/pytorch.pyi +2 -2
  161. metaflow-stubs/plugins/gcp/__init__.pyi +2 -2
  162. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +2 -2
  163. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +2 -2
  164. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +2 -2
  165. metaflow-stubs/plugins/gcp/gs_utils.pyi +2 -2
  166. metaflow-stubs/plugins/gcp/includefile_support.pyi +2 -2
  167. metaflow-stubs/plugins/kubernetes/__init__.pyi +2 -2
  168. metaflow-stubs/plugins/kubernetes/kube_utils.pyi +2 -2
  169. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +2 -2
  170. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +2 -2
  171. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +2 -2
  172. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +2 -2
  173. metaflow-stubs/plugins/kubernetes/spot_monitor_sidecar.pyi +2 -2
  174. metaflow-stubs/plugins/ollama/__init__.pyi +2 -2
  175. metaflow-stubs/plugins/parallel_decorator.pyi +2 -2
  176. metaflow-stubs/plugins/perimeters.pyi +2 -2
  177. metaflow-stubs/plugins/project_decorator.pyi +2 -2
  178. metaflow-stubs/plugins/pypi/__init__.pyi +3 -3
  179. metaflow-stubs/plugins/pypi/conda_decorator.pyi +2 -2
  180. metaflow-stubs/plugins/pypi/conda_environment.pyi +4 -4
  181. metaflow-stubs/plugins/pypi/parsers.pyi +2 -2
  182. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +2 -2
  183. metaflow-stubs/plugins/pypi/pypi_environment.pyi +2 -2
  184. metaflow-stubs/plugins/pypi/utils.pyi +2 -2
  185. metaflow-stubs/plugins/resources_decorator.pyi +2 -2
  186. metaflow-stubs/plugins/retry_decorator.pyi +2 -2
  187. metaflow-stubs/plugins/secrets/__init__.pyi +3 -3
  188. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +2 -2
  189. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +2 -2
  190. metaflow-stubs/plugins/snowflake/__init__.pyi +2 -2
  191. metaflow-stubs/plugins/storage_executor.pyi +2 -2
  192. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +3 -3
  193. metaflow-stubs/plugins/timeout_decorator.pyi +3 -3
  194. metaflow-stubs/plugins/torchtune/__init__.pyi +2 -2
  195. metaflow-stubs/plugins/uv/__init__.pyi +2 -2
  196. metaflow-stubs/plugins/uv/uv_environment.pyi +2 -2
  197. metaflow-stubs/profilers/__init__.pyi +2 -2
  198. metaflow-stubs/pylint_wrapper.pyi +2 -2
  199. metaflow-stubs/runner/__init__.pyi +2 -2
  200. metaflow-stubs/runner/deployer.pyi +29 -29
  201. metaflow-stubs/runner/deployer_impl.pyi +2 -2
  202. metaflow-stubs/runner/metaflow_runner.pyi +4 -4
  203. metaflow-stubs/runner/nbdeploy.pyi +2 -2
  204. metaflow-stubs/runner/nbrun.pyi +2 -2
  205. metaflow-stubs/runner/subprocess_manager.pyi +2 -2
  206. metaflow-stubs/runner/utils.pyi +2 -2
  207. metaflow-stubs/system/__init__.pyi +2 -2
  208. metaflow-stubs/system/system_logger.pyi +2 -2
  209. metaflow-stubs/system/system_monitor.pyi +2 -2
  210. metaflow-stubs/tagging_util.pyi +2 -2
  211. metaflow-stubs/tuple_util.pyi +2 -2
  212. metaflow-stubs/user_configs/__init__.pyi +2 -2
  213. metaflow-stubs/user_configs/config_decorators.pyi +6 -6
  214. metaflow-stubs/user_configs/config_options.pyi +2 -2
  215. metaflow-stubs/user_configs/config_parameters.pyi +7 -7
  216. {ob_metaflow_stubs-6.0.3.182rc2.dist-info → ob_metaflow_stubs-6.0.3.183rc1.dist-info}/METADATA +1 -1
  217. ob_metaflow_stubs-6.0.3.183rc1.dist-info/RECORD +220 -0
  218. ob_metaflow_stubs-6.0.3.182rc2.dist-info/RECORD +0 -220
  219. {ob_metaflow_stubs-6.0.3.182rc2.dist-info → ob_metaflow_stubs-6.0.3.183rc1.dist-info}/WHEEL +0 -0
  220. {ob_metaflow_stubs-6.0.3.182rc2.dist-info → ob_metaflow_stubs-6.0.3.183rc1.dist-info}/top_level.txt +0 -0
@@ -1,7 +1,7 @@
1
1
  ######################################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
- # MF version: 2.15.17.1+obcheckpoint(0.2.1);ob(v1) #
4
- # Generated on 2025-06-18T10:30:41.738497 #
3
+ # MF version: 2.15.18.1+obcheckpoint(0.2.1);ob(v1) #
4
+ # Generated on 2025-06-23T21:15:44.995494 #
5
5
  ######################################################################################################
6
6
 
7
7
  from __future__ import annotations
@@ -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 tuple_util as tuple_util
39
- from . import cards as cards
40
38
  from . import events as events
39
+ from . import cards as cards
41
40
  from . import metaflow_git as metaflow_git
41
+ from . import tuple_util as tuple_util
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
47
  from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
49
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
@@ -155,206 +155,38 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
155
155
  """
156
156
  ...
157
157
 
158
- 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]]]:
159
- """
160
- Specifies that this step is used to deploy an instance of the app.
161
- Requires that self.app_name, self.app_port, self.entrypoint and self.deployDir is set.
162
-
163
-
164
- Parameters
165
- ----------
166
- app_port : int
167
- Number of GPUs to use.
168
- app_name : str
169
- Name of the app to deploy.
170
- """
171
- ...
172
-
173
- @typing.overload
174
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
175
- """
176
- Decorator prototype for all step decorators. This function gets specialized
177
- and imported for all decorators types by _import_plugin_decorators().
178
- """
179
- ...
180
-
181
- @typing.overload
182
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
183
- ...
184
-
185
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
186
- """
187
- Decorator prototype for all step decorators. This function gets specialized
188
- and imported for all decorators types by _import_plugin_decorators().
189
- """
190
- ...
191
-
192
- @typing.overload
193
- 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]]]:
194
- """
195
- Enables loading / saving of models within a step.
196
-
197
-
198
-
199
- Parameters
200
- ----------
201
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
202
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
203
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
204
- - `current.checkpoint`
205
- - `current.model`
206
- - `current.huggingface_hub`
207
-
208
- 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
209
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
210
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
211
-
212
- temp_dir_root : str, default: None
213
- The root directory under which `current.model.loaded` will store loaded models
214
- """
215
- ...
216
-
217
- @typing.overload
218
- def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
219
- ...
220
-
221
- @typing.overload
222
- def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
223
- ...
224
-
225
- 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):
226
- """
227
- Enables loading / saving of models within a step.
228
-
229
-
230
-
231
- Parameters
232
- ----------
233
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
234
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
235
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
236
- - `current.checkpoint`
237
- - `current.model`
238
- - `current.huggingface_hub`
239
-
240
- 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
241
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
242
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
243
-
244
- temp_dir_root : str, default: None
245
- The root directory under which `current.model.loaded` will store loaded models
246
- """
247
- ...
248
-
249
- 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]]]:
250
- """
251
- Specifies that this step should execute on DGX cloud.
252
-
253
-
254
- Parameters
255
- ----------
256
- gpu : int
257
- Number of GPUs to use.
258
- gpu_type : str
259
- Type of Nvidia GPU to use.
260
- queue_timeout : int
261
- Time to keep the job in NVCF's queue.
262
- """
263
- ...
264
-
265
158
  @typing.overload
266
- 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]]]:
159
+ 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]]]:
267
160
  """
268
- Enables checkpointing for a step.
269
-
161
+ Specifies secrets to be retrieved and injected as environment variables prior to
162
+ the execution of a step.
270
163
 
271
164
 
272
165
  Parameters
273
166
  ----------
274
- load_policy : str, default: "fresh"
275
- The policy for loading the checkpoint. The following policies are supported:
276
- - "eager": Loads the the latest available checkpoint within the namespace.
277
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
278
- will be loaded at the start of the task.
279
- - "none": Do not load any checkpoint
280
- - "fresh": Loads the lastest checkpoint created within the running Task.
281
- This mode helps loading checkpoints across various retry attempts of the same task.
282
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
283
- created within the task will be loaded when the task is retries execution on failure.
284
-
285
- temp_dir_root : str, default: None
286
- The root directory under which `current.checkpoint.directory` will be created.
167
+ sources : List[Union[str, Dict[str, Any]]], default: []
168
+ List of secret specs, defining how the secrets are to be retrieved
287
169
  """
288
170
  ...
289
171
 
290
172
  @typing.overload
291
- def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
173
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
292
174
  ...
293
175
 
294
176
  @typing.overload
295
- def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
296
- ...
297
-
298
- 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):
299
- """
300
- Enables checkpointing for a step.
301
-
302
-
303
-
304
- Parameters
305
- ----------
306
- load_policy : str, default: "fresh"
307
- The policy for loading the checkpoint. The following policies are supported:
308
- - "eager": Loads the the latest available checkpoint within the namespace.
309
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
310
- will be loaded at the start of the task.
311
- - "none": Do not load any checkpoint
312
- - "fresh": Loads the lastest checkpoint created within the running Task.
313
- This mode helps loading checkpoints across various retry attempts of the same task.
314
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
315
- created within the task will be loaded when the task is retries execution on failure.
316
-
317
- temp_dir_root : str, default: None
318
- The root directory under which `current.checkpoint.directory` will be created.
319
- """
320
- ...
321
-
322
- def nvct(*, gpu: int, gpu_type: str) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
323
- """
324
- Specifies that this step should execute on DGX cloud.
325
-
326
-
327
- Parameters
328
- ----------
329
- gpu : int
330
- Number of GPUs to use.
331
- gpu_type : str
332
- Type of Nvidia GPU to use.
333
- """
177
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
334
178
  ...
335
179
 
336
- 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]]]:
180
+ 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]]] = []):
337
181
  """
338
- Decorator that helps cache, version and store models/datasets from huggingface hub.
182
+ Specifies secrets to be retrieved and injected as environment variables prior to
183
+ the execution of a step.
339
184
 
340
185
 
341
186
  Parameters
342
187
  ----------
343
- temp_dir_root : str, optional
344
- The root directory that will hold the temporary directory where objects will be downloaded.
345
-
346
- load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
347
- The list of repos (models/datasets) to load.
348
-
349
- Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
350
-
351
- - If repo (model/dataset) is not found in the datastore:
352
- - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
353
- - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
354
- - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
355
-
356
- - If repo is found in the datastore:
357
- - Loads it directly from datastore to local path (can be temporary directory or specified path)
188
+ sources : List[Union[str, Dict[str, Any]]], default: []
189
+ List of secret specs, defining how the secrets are to be retrieved
358
190
  """
359
191
  ...
360
192
 
@@ -507,110 +339,55 @@ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: ty
507
339
  ...
508
340
 
509
341
  @typing.overload
510
- 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]]]:
342
+ def resources(*, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
511
343
  """
512
- Specifies the number of times the task corresponding
513
- to a step needs to be retried.
344
+ Specifies the resources needed when executing this step.
514
345
 
515
- This decorator is useful for handling transient errors, such as networking issues.
516
- If your task contains operations that can't be retried safely, e.g. database updates,
517
- it is advisable to annotate it with `@retry(times=0)`.
346
+ Use `@resources` to specify the resource requirements
347
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
518
348
 
519
- This can be used in conjunction with the `@catch` decorator. The `@catch`
520
- decorator will execute a no-op task after all retries have been exhausted,
521
- ensuring that the flow execution can continue.
349
+ You can choose the compute layer on the command line by executing e.g.
350
+ ```
351
+ python myflow.py run --with batch
352
+ ```
353
+ or
354
+ ```
355
+ python myflow.py run --with kubernetes
356
+ ```
357
+ which executes the flow on the desired system using the
358
+ requirements specified in `@resources`.
522
359
 
523
360
 
524
361
  Parameters
525
362
  ----------
526
- times : int, default 3
527
- Number of times to retry this task.
528
- minutes_between_retries : int, default 2
529
- Number of minutes between retries.
363
+ cpu : int, default 1
364
+ Number of CPUs required for this step.
365
+ gpu : int, optional, default None
366
+ Number of GPUs required for this step.
367
+ disk : int, optional, default None
368
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
369
+ memory : int, default 4096
370
+ Memory size (in MB) required for this step.
371
+ shared_memory : int, optional, default None
372
+ The value for the size (in MiB) of the /dev/shm volume for this step.
373
+ This parameter maps to the `--shm-size` option in Docker.
530
374
  """
531
375
  ...
532
376
 
533
377
  @typing.overload
534
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
378
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
535
379
  ...
536
380
 
537
381
  @typing.overload
538
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
382
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
539
383
  ...
540
384
 
541
- 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):
385
+ def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None):
542
386
  """
543
- Specifies the number of times the task corresponding
544
- to a step needs to be retried.
545
-
546
- This decorator is useful for handling transient errors, such as networking issues.
547
- If your task contains operations that can't be retried safely, e.g. database updates,
548
- it is advisable to annotate it with `@retry(times=0)`.
549
-
550
- This can be used in conjunction with the `@catch` decorator. The `@catch`
551
- decorator will execute a no-op task after all retries have been exhausted,
552
- ensuring that the flow execution can continue.
387
+ Specifies the resources needed when executing this step.
553
388
 
554
-
555
- Parameters
556
- ----------
557
- times : int, default 3
558
- Number of times to retry this task.
559
- minutes_between_retries : int, default 2
560
- Number of minutes between retries.
561
- """
562
- ...
563
-
564
- def ollama(*, models: list, backend: str, force_pull: bool, cache_update_policy: str, force_cache_update: bool, debug: bool, circuit_breaker_config: dict, timeout_config: dict) -> 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]]]:
565
- """
566
- This decorator is used to run Ollama APIs as Metaflow task sidecars.
567
-
568
- User code call
569
- --------------
570
- @ollama(
571
- models=[...],
572
- ...
573
- )
574
-
575
- Valid backend options
576
- ---------------------
577
- - 'local': Run as a separate process on the local task machine.
578
- - (TODO) 'managed': Outerbounds hosts and selects compute provider.
579
- - (TODO) 'remote': Spin up separate instance to serve Ollama models.
580
-
581
- Valid model options
582
- -------------------
583
- Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
584
-
585
-
586
- Parameters
587
- ----------
588
- models: list[str]
589
- List of Ollama containers running models in sidecars.
590
- backend: str
591
- Determines where and how to run the Ollama process.
592
- force_pull: bool
593
- Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
594
- cache_update_policy: str
595
- Cache update policy: "auto", "force", or "never".
596
- force_cache_update: bool
597
- Simple override for "force" cache update policy.
598
- debug: bool
599
- Whether to turn on verbose debugging logs.
600
- circuit_breaker_config: dict
601
- Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
602
- timeout_config: dict
603
- Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
604
- """
605
- ...
606
-
607
- @typing.overload
608
- def resources(*, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
609
- """
610
- Specifies the resources needed when executing this step.
611
-
612
- Use `@resources` to specify the resource requirements
613
- independently of the specific compute layer (`@batch`, `@kubernetes`).
389
+ Use `@resources` to specify the resource requirements
390
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
614
391
 
615
392
  You can choose the compute layer on the command line by executing e.g.
616
393
  ```
@@ -640,46 +417,55 @@ def resources(*, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Op
640
417
  """
641
418
  ...
642
419
 
643
- @typing.overload
644
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
645
- ...
646
-
647
- @typing.overload
648
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
420
+ def nvct(*, gpu: int, gpu_type: str) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
421
+ """
422
+ Specifies that this step should execute on DGX cloud.
423
+
424
+
425
+ Parameters
426
+ ----------
427
+ gpu : int
428
+ Number of GPUs to use.
429
+ gpu_type : str
430
+ Type of Nvidia GPU to use.
431
+ """
649
432
  ...
650
433
 
651
- def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None):
434
+ def vllm(*, model: str, backend: str, debug: bool, kwargs: 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]]]:
652
435
  """
653
- Specifies the resources needed when executing this step.
436
+ This decorator is used to run vllm APIs as Metaflow task sidecars.
654
437
 
655
- Use `@resources` to specify the resource requirements
656
- independently of the specific compute layer (`@batch`, `@kubernetes`).
438
+ User code call
439
+ --------------
440
+ @vllm(
441
+ model="...",
442
+ ...
443
+ )
657
444
 
658
- You can choose the compute layer on the command line by executing e.g.
659
- ```
660
- python myflow.py run --with batch
661
- ```
662
- or
663
- ```
664
- python myflow.py run --with kubernetes
665
- ```
666
- which executes the flow on the desired system using the
667
- requirements specified in `@resources`.
445
+ Valid backend options
446
+ ---------------------
447
+ - 'local': Run as a separate process on the local task machine.
448
+
449
+ Valid model options
450
+ -------------------
451
+ Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
452
+
453
+ NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
454
+ If you need multiple models, you must create multiple @vllm decorators.
668
455
 
669
456
 
670
457
  Parameters
671
458
  ----------
672
- cpu : int, default 1
673
- Number of CPUs required for this step.
674
- gpu : int, optional, default None
675
- Number of GPUs required for this step.
676
- disk : int, optional, default None
677
- Disk size (in MB) required for this step. Only applies on Kubernetes.
678
- memory : int, default 4096
679
- Memory size (in MB) required for this step.
680
- shared_memory : int, optional, default None
681
- The value for the size (in MiB) of the /dev/shm volume for this step.
682
- This parameter maps to the `--shm-size` option in Docker.
459
+ model: str
460
+ HuggingFace model identifier to be served by vLLM.
461
+ backend: str
462
+ Determines where and how to run the vLLM process.
463
+ debug: bool
464
+ Whether to turn on verbose debugging logs.
465
+ kwargs : Any
466
+ Any other keyword arguments are passed directly to the vLLM engine.
467
+ This allows for flexible configuration of vLLM server settings.
468
+ For example, `tensor_parallel_size=2`.
683
469
  """
684
470
  ...
685
471
 
@@ -735,174 +521,302 @@ def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
735
521
  ...
736
522
 
737
523
  @typing.overload
738
- 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]]]:
524
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
739
525
  """
740
- Specifies secrets to be retrieved and injected as environment variables prior to
741
- the execution of a step.
526
+ Internal decorator to support Fast bakery
527
+ """
528
+ ...
529
+
530
+ @typing.overload
531
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
532
+ ...
533
+
534
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
535
+ """
536
+ Internal decorator to support Fast bakery
537
+ """
538
+ ...
539
+
540
+ @typing.overload
541
+ 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]]]:
542
+ """
543
+ Specifies the PyPI packages for the step.
544
+
545
+ Information in this decorator will augment any
546
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
547
+ you can use `@pypi_base` to set packages required by all
548
+ steps and use `@pypi` to specify step-specific overrides.
742
549
 
743
550
 
744
551
  Parameters
745
552
  ----------
746
- sources : List[Union[str, Dict[str, Any]]], default: []
747
- List of secret specs, defining how the secrets are to be retrieved
553
+ packages : Dict[str, str], default: {}
554
+ Packages to use for this step. The key is the name of the package
555
+ and the value is the version to use.
556
+ python : str, optional, default: None
557
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
558
+ that the version used will correspond to the version of the Python interpreter used to start the run.
748
559
  """
749
560
  ...
750
561
 
751
562
  @typing.overload
752
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
563
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
753
564
  ...
754
565
 
755
566
  @typing.overload
756
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
567
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
757
568
  ...
758
569
 
759
- 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]]] = []):
570
+ 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):
760
571
  """
761
- Specifies secrets to be retrieved and injected as environment variables prior to
762
- the execution of a step.
572
+ Specifies the PyPI packages for the step.
573
+
574
+ Information in this decorator will augment any
575
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
576
+ you can use `@pypi_base` to set packages required by all
577
+ steps and use `@pypi` to specify step-specific overrides.
763
578
 
764
579
 
765
580
  Parameters
766
581
  ----------
767
- sources : List[Union[str, Dict[str, Any]]], default: []
768
- List of secret specs, defining how the secrets are to be retrieved
582
+ packages : Dict[str, str], default: {}
583
+ Packages to use for this step. The key is the name of the package
584
+ and the value is the version to use.
585
+ python : str, optional, default: None
586
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
587
+ that the version used will correspond to the version of the Python interpreter used to start the run.
769
588
  """
770
589
  ...
771
590
 
772
- @typing.overload
773
- 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]]]:
591
+ 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]]]:
774
592
  """
775
- Specifies environment variables to be set prior to the execution of a step.
593
+ Specifies that this step should execute on DGX cloud.
776
594
 
777
595
 
778
596
  Parameters
779
597
  ----------
780
- vars : Dict[str, str], default {}
781
- Dictionary of environment variables to set.
598
+ gpu : int
599
+ Number of GPUs to use.
600
+ gpu_type : str
601
+ Type of Nvidia GPU to use.
602
+ queue_timeout : int
603
+ Time to keep the job in NVCF's queue.
782
604
  """
783
605
  ...
784
606
 
785
- @typing.overload
786
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
787
- ...
788
-
789
- @typing.overload
790
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
791
- ...
792
-
793
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
607
+ 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]]]:
794
608
  """
795
- Specifies environment variables to be set prior to the execution of a step.
609
+ Specifies that this step is used to deploy an instance of the app.
610
+ Requires that self.app_name, self.app_port, self.entrypoint and self.deployDir is set.
796
611
 
797
612
 
798
613
  Parameters
799
614
  ----------
800
- vars : Dict[str, str], default {}
801
- Dictionary of environment variables to set.
615
+ app_port : int
616
+ Number of GPUs to use.
617
+ app_name : str
618
+ Name of the app to deploy.
802
619
  """
803
620
  ...
804
621
 
805
622
  @typing.overload
806
- 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]]]:
623
+ 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]]]:
807
624
  """
808
- Creates a human-readable report, a Metaflow Card, after this step completes.
625
+ Enables loading / saving of models within a step.
809
626
 
810
- Note that you may add multiple `@card` decorators in a step with different parameters.
811
627
 
812
628
 
813
629
  Parameters
814
630
  ----------
815
- type : str, default 'default'
816
- Card type.
817
- id : str, optional, default None
818
- If multiple cards are present, use this id to identify this card.
819
- options : Dict[str, Any], default {}
820
- Options passed to the card. The contents depend on the card type.
821
- timeout : int, default 45
822
- Interrupt reporting if it takes more than this many seconds.
631
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
632
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
633
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
634
+ - `current.checkpoint`
635
+ - `current.model`
636
+ - `current.huggingface_hub`
637
+
638
+ 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
639
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
640
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
641
+
642
+ temp_dir_root : str, default: None
643
+ The root directory under which `current.model.loaded` will store loaded models
823
644
  """
824
645
  ...
825
646
 
826
647
  @typing.overload
827
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
648
+ def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
828
649
  ...
829
650
 
830
651
  @typing.overload
831
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
652
+ def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
832
653
  ...
833
654
 
834
- 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):
655
+ 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):
835
656
  """
836
- Creates a human-readable report, a Metaflow Card, after this step completes.
657
+ Enables loading / saving of models within a step.
837
658
 
838
- Note that you may add multiple `@card` decorators in a step with different parameters.
839
659
 
840
660
 
841
661
  Parameters
842
662
  ----------
843
- type : str, default 'default'
844
- Card type.
845
- id : str, optional, default None
846
- If multiple cards are present, use this id to identify this card.
847
- options : Dict[str, Any], default {}
848
- Options passed to the card. The contents depend on the card type.
849
- timeout : int, default 45
850
- Interrupt reporting if it takes more than this many seconds.
663
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
664
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
665
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
666
+ - `current.checkpoint`
667
+ - `current.model`
668
+ - `current.huggingface_hub`
669
+
670
+ 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
671
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
672
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
673
+
674
+ temp_dir_root : str, default: None
675
+ The root directory under which `current.model.loaded` will store loaded models
851
676
  """
852
677
  ...
853
678
 
854
679
  @typing.overload
855
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
680
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
856
681
  """
857
- Internal decorator to support Fast bakery
682
+ Decorator prototype for all step decorators. This function gets specialized
683
+ and imported for all decorators types by _import_plugin_decorators().
858
684
  """
859
685
  ...
860
686
 
861
687
  @typing.overload
862
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
688
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
863
689
  ...
864
690
 
865
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
691
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
866
692
  """
867
- Internal decorator to support Fast bakery
693
+ Decorator prototype for all step decorators. This function gets specialized
694
+ and imported for all decorators types by _import_plugin_decorators().
868
695
  """
869
696
  ...
870
697
 
871
- def vllm(*, model: str, backend: str, debug: bool, kwargs: 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]]]:
698
+ 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]]]:
872
699
  """
873
- This decorator is used to run vllm APIs as Metaflow task sidecars.
700
+ Decorator that helps cache, version and store models/datasets from huggingface hub.
701
+
702
+
703
+ Parameters
704
+ ----------
705
+ temp_dir_root : str, optional
706
+ The root directory that will hold the temporary directory where objects will be downloaded.
707
+
708
+ load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
709
+ The list of repos (models/datasets) to load.
710
+
711
+ Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
712
+
713
+ - If repo (model/dataset) is not found in the datastore:
714
+ - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
715
+ - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
716
+ - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
717
+
718
+ - If repo is found in the datastore:
719
+ - Loads it directly from datastore to local path (can be temporary directory or specified path)
720
+ """
721
+ ...
722
+
723
+ @typing.overload
724
+ 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]]]:
725
+ """
726
+ Enables checkpointing for a step.
727
+
728
+
729
+
730
+ Parameters
731
+ ----------
732
+ load_policy : str, default: "fresh"
733
+ The policy for loading the checkpoint. The following policies are supported:
734
+ - "eager": Loads the the latest available checkpoint within the namespace.
735
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
736
+ will be loaded at the start of the task.
737
+ - "none": Do not load any checkpoint
738
+ - "fresh": Loads the lastest checkpoint created within the running Task.
739
+ This mode helps loading checkpoints across various retry attempts of the same task.
740
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
741
+ created within the task will be loaded when the task is retries execution on failure.
742
+
743
+ temp_dir_root : str, default: None
744
+ The root directory under which `current.checkpoint.directory` will be created.
745
+ """
746
+ ...
747
+
748
+ @typing.overload
749
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
750
+ ...
751
+
752
+ @typing.overload
753
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
754
+ ...
755
+
756
+ 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):
757
+ """
758
+ Enables checkpointing for a step.
759
+
760
+
761
+
762
+ Parameters
763
+ ----------
764
+ load_policy : str, default: "fresh"
765
+ The policy for loading the checkpoint. The following policies are supported:
766
+ - "eager": Loads the the latest available checkpoint within the namespace.
767
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
768
+ will be loaded at the start of the task.
769
+ - "none": Do not load any checkpoint
770
+ - "fresh": Loads the lastest checkpoint created within the running Task.
771
+ This mode helps loading checkpoints across various retry attempts of the same task.
772
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
773
+ created within the task will be loaded when the task is retries execution on failure.
774
+
775
+ temp_dir_root : str, default: None
776
+ The root directory under which `current.checkpoint.directory` will be created.
777
+ """
778
+ ...
779
+
780
+ def ollama(*, models: list, backend: str, force_pull: bool, cache_update_policy: str, force_cache_update: bool, debug: bool, circuit_breaker_config: dict, timeout_config: dict) -> 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]]]:
781
+ """
782
+ This decorator is used to run Ollama APIs as Metaflow task sidecars.
874
783
 
875
784
  User code call
876
785
  --------------
877
- @vllm(
878
- model="...",
786
+ @ollama(
787
+ models=[...],
879
788
  ...
880
789
  )
881
790
 
882
791
  Valid backend options
883
792
  ---------------------
884
793
  - 'local': Run as a separate process on the local task machine.
794
+ - (TODO) 'managed': Outerbounds hosts and selects compute provider.
795
+ - (TODO) 'remote': Spin up separate instance to serve Ollama models.
885
796
 
886
797
  Valid model options
887
798
  -------------------
888
- Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
889
-
890
- NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
891
- If you need multiple models, you must create multiple @vllm decorators.
799
+ Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
892
800
 
893
801
 
894
802
  Parameters
895
803
  ----------
896
- model: str
897
- HuggingFace model identifier to be served by vLLM.
804
+ models: list[str]
805
+ List of Ollama containers running models in sidecars.
898
806
  backend: str
899
- Determines where and how to run the vLLM process.
807
+ Determines where and how to run the Ollama process.
808
+ force_pull: bool
809
+ Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
810
+ cache_update_policy: str
811
+ Cache update policy: "auto", "force", or "never".
812
+ force_cache_update: bool
813
+ Simple override for "force" cache update policy.
900
814
  debug: bool
901
815
  Whether to turn on verbose debugging logs.
902
- kwargs : Any
903
- Any other keyword arguments are passed directly to the vLLM engine.
904
- This allows for flexible configuration of vLLM server settings.
905
- For example, `tensor_parallel_size=2`.
816
+ circuit_breaker_config: dict
817
+ Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
818
+ timeout_config: dict
819
+ Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
906
820
  """
907
821
  ...
908
822
 
@@ -966,223 +880,139 @@ def conda(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
966
880
  ...
967
881
 
968
882
  @typing.overload
969
- 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]]]:
883
+ 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]]]:
970
884
  """
971
- Specifies the PyPI packages for the step.
972
-
973
- Information in this decorator will augment any
974
- attributes set in the `@pyi_base` flow-level decorator. Hence,
975
- you can use `@pypi_base` to set packages required by all
976
- steps and use `@pypi` to specify step-specific overrides.
885
+ Specifies environment variables to be set prior to the execution of a step.
977
886
 
978
887
 
979
888
  Parameters
980
889
  ----------
981
- packages : Dict[str, str], default: {}
982
- Packages to use for this step. The key is the name of the package
983
- and the value is the version to use.
984
- python : str, optional, default: None
985
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
986
- that the version used will correspond to the version of the Python interpreter used to start the run.
890
+ vars : Dict[str, str], default {}
891
+ Dictionary of environment variables to set.
987
892
  """
988
893
  ...
989
894
 
990
895
  @typing.overload
991
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
896
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
992
897
  ...
993
898
 
994
899
  @typing.overload
995
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
900
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
996
901
  ...
997
902
 
998
- 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):
903
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
999
904
  """
1000
- Specifies the PyPI packages for the step.
1001
-
1002
- Information in this decorator will augment any
1003
- attributes set in the `@pyi_base` flow-level decorator. Hence,
1004
- you can use `@pypi_base` to set packages required by all
1005
- steps and use `@pypi` to specify step-specific overrides.
905
+ Specifies environment variables to be set prior to the execution of a step.
1006
906
 
1007
907
 
1008
908
  Parameters
1009
909
  ----------
1010
- packages : Dict[str, str], default: {}
1011
- Packages to use for this step. The key is the name of the package
1012
- and the value is the version to use.
1013
- python : str, optional, default: None
1014
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1015
- that the version used will correspond to the version of the Python interpreter used to start the run.
910
+ vars : Dict[str, str], default {}
911
+ Dictionary of environment variables to set.
1016
912
  """
1017
913
  ...
1018
914
 
1019
- def airflow_external_task_sensor(*, timeout: int, poke_interval: int, mode: str, exponential_backoff: bool, pool: str, soft_fail: bool, name: str, description: str, external_dag_id: str, external_task_ids: typing.List[str], allowed_states: typing.List[str], failed_states: typing.List[str], execution_delta: "datetime.timedelta", check_existence: bool) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
915
+ @typing.overload
916
+ 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]]]:
1020
917
  """
1021
- The `@airflow_external_task_sensor` decorator attaches a Airflow [ExternalTaskSensor](https://airflow.apache.org/docs/apache-airflow/stable/_api/airflow/sensors/external_task/index.html#airflow.sensors.external_task.ExternalTaskSensor) before the start step of the flow.
1022
- This decorator only works when a flow is scheduled on Airflow and is compiled using `airflow create`. More than one `@airflow_external_task_sensor` can be added as a flow decorators. Adding more than one decorator will ensure that `start` step starts only after all sensors finish.
918
+ Creates a human-readable report, a Metaflow Card, after this step completes.
919
+
920
+ Note that you may add multiple `@card` decorators in a step with different parameters.
1023
921
 
1024
922
 
1025
923
  Parameters
1026
924
  ----------
1027
- timeout : int
1028
- Time, in seconds before the task times out and fails. (Default: 3600)
1029
- poke_interval : int
1030
- Time in seconds that the job should wait in between each try. (Default: 60)
1031
- mode : str
1032
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1033
- exponential_backoff : bool
1034
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1035
- pool : str
1036
- the slot pool this task should run in,
1037
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1038
- soft_fail : bool
1039
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1040
- name : str
1041
- Name of the sensor on Airflow
1042
- description : str
1043
- Description of sensor in the Airflow UI
1044
- external_dag_id : str
1045
- The dag_id that contains the task you want to wait for.
1046
- external_task_ids : List[str]
1047
- The list of task_ids that you want to wait for.
1048
- If None (default value) the sensor waits for the DAG. (Default: None)
1049
- allowed_states : List[str]
1050
- Iterable of allowed states, (Default: ['success'])
1051
- failed_states : List[str]
1052
- Iterable of failed or dis-allowed states. (Default: None)
1053
- execution_delta : datetime.timedelta
1054
- time difference with the previous execution to look at,
1055
- the default is the same logical date as the current task or DAG. (Default: None)
1056
- check_existence: bool
1057
- Set to True to check if the external task exists or check if
1058
- the DAG to wait for exists. (Default: True)
925
+ type : str, default 'default'
926
+ Card type.
927
+ id : str, optional, default None
928
+ If multiple cards are present, use this id to identify this card.
929
+ options : Dict[str, Any], default {}
930
+ Options passed to the card. The contents depend on the card type.
931
+ timeout : int, default 45
932
+ Interrupt reporting if it takes more than this many seconds.
1059
933
  """
1060
934
  ...
1061
935
 
1062
936
  @typing.overload
1063
- def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1064
- """
1065
- Specifies the PyPI packages for all steps of the flow.
1066
-
1067
- Use `@pypi_base` to set common packages required by all
1068
- steps and use `@pypi` to specify step-specific overrides.
1069
-
1070
- Parameters
1071
- ----------
1072
- packages : Dict[str, str], default: {}
1073
- Packages to use for this flow. The key is the name of the package
1074
- and the value is the version to use.
1075
- python : str, optional, default: None
1076
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1077
- that the version used will correspond to the version of the Python interpreter used to start the run.
1078
- """
937
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1079
938
  ...
1080
939
 
1081
940
  @typing.overload
1082
- def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
941
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1083
942
  ...
1084
943
 
1085
- def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
944
+ 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):
1086
945
  """
1087
- Specifies the PyPI packages for all steps of the flow.
946
+ Creates a human-readable report, a Metaflow Card, after this step completes.
947
+
948
+ Note that you may add multiple `@card` decorators in a step with different parameters.
1088
949
 
1089
- Use `@pypi_base` to set common packages required by all
1090
- steps and use `@pypi` to specify step-specific overrides.
1091
950
 
1092
951
  Parameters
1093
952
  ----------
1094
- packages : Dict[str, str], default: {}
1095
- Packages to use for this flow. The key is the name of the package
1096
- and the value is the version to use.
1097
- python : str, optional, default: None
1098
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1099
- that the version used will correspond to the version of the Python interpreter used to start the run.
953
+ type : str, default 'default'
954
+ Card type.
955
+ id : str, optional, default None
956
+ If multiple cards are present, use this id to identify this card.
957
+ options : Dict[str, Any], default {}
958
+ Options passed to the card. The contents depend on the card type.
959
+ timeout : int, default 45
960
+ Interrupt reporting if it takes more than this many seconds.
1100
961
  """
1101
962
  ...
1102
963
 
1103
- def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
964
+ @typing.overload
965
+ 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]]]:
1104
966
  """
1105
- Specifies what flows belong to the same project.
967
+ Specifies the number of times the task corresponding
968
+ to a step needs to be retried.
1106
969
 
1107
- A project-specific namespace is created for all flows that
1108
- use the same `@project(name)`.
970
+ This decorator is useful for handling transient errors, such as networking issues.
971
+ If your task contains operations that can't be retried safely, e.g. database updates,
972
+ it is advisable to annotate it with `@retry(times=0)`.
973
+
974
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
975
+ decorator will execute a no-op task after all retries have been exhausted,
976
+ ensuring that the flow execution can continue.
1109
977
 
1110
978
 
1111
979
  Parameters
1112
980
  ----------
1113
- name : str
1114
- Project name. Make sure that the name is unique amongst all
1115
- projects that use the same production scheduler. The name may
1116
- contain only lowercase alphanumeric characters and underscores.
1117
-
1118
- branch : Optional[str], default None
1119
- The branch to use. If not specified, the branch is set to
1120
- `user.<username>` unless `production` is set to `True`. This can
1121
- also be set on the command line using `--branch` as a top-level option.
1122
- It is an error to specify `branch` in the decorator and on the command line.
1123
-
1124
- production : bool, default False
1125
- Whether or not the branch is the production branch. This can also be set on the
1126
- command line using `--production` as a top-level option. It is an error to specify
1127
- `production` in the decorator and on the command line.
1128
- The project branch name will be:
1129
- - if `branch` is specified:
1130
- - if `production` is True: `prod.<branch>`
1131
- - if `production` is False: `test.<branch>`
1132
- - if `branch` is not specified:
1133
- - if `production` is True: `prod`
1134
- - if `production` is False: `user.<username>`
981
+ times : int, default 3
982
+ Number of times to retry this task.
983
+ minutes_between_retries : int, default 2
984
+ Number of minutes between retries.
1135
985
  """
1136
986
  ...
1137
987
 
1138
988
  @typing.overload
1139
- def schedule(*, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1140
- """
1141
- Specifies the times when the flow should be run when running on a
1142
- production scheduler.
1143
-
1144
-
1145
- Parameters
1146
- ----------
1147
- hourly : bool, default False
1148
- Run the workflow hourly.
1149
- daily : bool, default True
1150
- Run the workflow daily.
1151
- weekly : bool, default False
1152
- Run the workflow weekly.
1153
- cron : str, optional, default None
1154
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1155
- specified by this expression.
1156
- timezone : str, optional, default None
1157
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1158
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1159
- """
989
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1160
990
  ...
1161
991
 
1162
992
  @typing.overload
1163
- def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
993
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1164
994
  ...
1165
995
 
1166
- def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None):
996
+ 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):
1167
997
  """
1168
- Specifies the times when the flow should be run when running on a
1169
- production scheduler.
998
+ Specifies the number of times the task corresponding
999
+ to a step needs to be retried.
1000
+
1001
+ This decorator is useful for handling transient errors, such as networking issues.
1002
+ If your task contains operations that can't be retried safely, e.g. database updates,
1003
+ it is advisable to annotate it with `@retry(times=0)`.
1004
+
1005
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
1006
+ decorator will execute a no-op task after all retries have been exhausted,
1007
+ ensuring that the flow execution can continue.
1170
1008
 
1171
1009
 
1172
1010
  Parameters
1173
1011
  ----------
1174
- hourly : bool, default False
1175
- Run the workflow hourly.
1176
- daily : bool, default True
1177
- Run the workflow daily.
1178
- weekly : bool, default False
1179
- Run the workflow weekly.
1180
- cron : str, optional, default None
1181
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1182
- specified by this expression.
1183
- timezone : str, optional, default None
1184
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1185
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1012
+ times : int, default 3
1013
+ Number of times to retry this task.
1014
+ minutes_between_retries : int, default 2
1015
+ Number of minutes between retries.
1186
1016
  """
1187
1017
  ...
1188
1018
 
@@ -1229,151 +1059,138 @@ def airflow_s3_key_sensor(*, timeout: int, poke_interval: int, mode: str, expone
1229
1059
  """
1230
1060
  ...
1231
1061
 
1232
- def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None):
1233
- """
1234
- Allows setting external datastores to save data for the
1235
- `@checkpoint`/`@model`/`@huggingface_hub` decorators.
1236
-
1237
- This decorator is useful when users wish to save data to a different datastore
1238
- than what is configured in Metaflow. This can be for variety of reasons:
1239
-
1240
- 1. Data security: The objects needs to be stored in a bucket (object storage) that is not accessible by other flows.
1241
- 2. Data Locality: The location where the task is executing is not located in the same region as the datastore.
1242
- - Example: Metaflow datastore lives in US East, but the task is executing in Finland datacenters.
1243
- 3. Data Lifecycle Policies: The objects need to be archived / managed separately from the Metaflow managed objects.
1244
- - Example: Flow is training very large models that need to be stored separately and will be deleted more aggressively than the Metaflow managed objects.
1245
-
1246
- Usage:
1247
- ----------
1248
-
1249
- - Using a custom IAM role to access the datastore.
1250
-
1251
- ```python
1252
- @with_artifact_store(
1253
- type="s3",
1254
- config=lambda: {
1255
- "root": "s3://my-bucket-foo/path/to/root",
1256
- "role_arn": ROLE,
1257
- },
1258
- )
1259
- class MyFlow(FlowSpec):
1260
-
1261
- @checkpoint
1262
- @step
1263
- def start(self):
1264
- with open("my_file.txt", "w") as f:
1265
- f.write("Hello, World!")
1266
- self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1267
- self.next(self.end)
1268
-
1269
- ```
1270
-
1271
- - Using credentials to access the s3-compatible datastore.
1272
-
1273
- ```python
1274
- @with_artifact_store(
1275
- type="s3",
1276
- config=lambda: {
1277
- "root": "s3://my-bucket-foo/path/to/root",
1278
- "client_params": {
1279
- "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1280
- "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1281
- },
1282
- },
1283
- )
1284
- class MyFlow(FlowSpec):
1285
-
1286
- @checkpoint
1287
- @step
1288
- def start(self):
1289
- with open("my_file.txt", "w") as f:
1290
- f.write("Hello, World!")
1291
- self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1292
- self.next(self.end)
1293
-
1294
- ```
1295
-
1296
- - Accessing objects stored in external datastores after task execution.
1297
-
1298
- ```python
1299
- run = Run("CheckpointsTestsFlow/8992")
1300
- with artifact_store_from(run=run, config={
1301
- "client_params": {
1302
- "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1303
- "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1304
- },
1305
- }):
1306
- with Checkpoint() as cp:
1307
- latest = cp.list(
1308
- task=run["start"].task
1309
- )[0]
1310
- print(latest)
1311
- cp.load(
1312
- latest,
1313
- "test-checkpoints"
1314
- )
1315
-
1316
- task = Task("TorchTuneFlow/8484/train/53673")
1317
- with artifact_store_from(run=run, config={
1318
- "client_params": {
1319
- "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1320
- "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1321
- },
1322
- }):
1323
- load_model(
1324
- task.data.model_ref,
1325
- "test-models"
1326
- )
1327
- ```
1328
- Parameters:
1329
- ----------
1330
-
1331
- type: str
1332
- The type of the datastore. Can be one of 's3', 'gcs', 'azure' or any other supported metaflow Datastore.
1333
-
1334
- config: dict or Callable
1335
- Dictionary of configuration options for the datastore. The following keys are required:
1336
- - root: The root path in the datastore where the data will be saved. (needs to be in the format expected by the datastore)
1337
- - example: 's3://bucket-name/path/to/root'
1338
- - example: 'gs://bucket-name/path/to/root'
1339
- - example: 'https://myblockacc.blob.core.windows.net/metaflow/'
1340
- - role_arn (optional): AWS IAM role to access s3 bucket (only when `type` is 's3')
1341
- - session_vars (optional): AWS session variables to access s3 bucket (only when `type` is 's3')
1342
- - client_params (optional): AWS client parameters to access s3 bucket (only when `type` is 's3')
1343
- """
1344
- ...
1345
-
1346
1062
  @typing.overload
1347
- 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]]:
1063
+ 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]]:
1348
1064
  """
1349
- Specifies the event(s) that this flow depends on.
1065
+ Specifies the flow(s) that this flow depends on.
1350
1066
 
1351
1067
  ```
1352
- @trigger(event='foo')
1068
+ @trigger_on_finish(flow='FooFlow')
1353
1069
  ```
1354
1070
  or
1355
1071
  ```
1356
- @trigger(events=['foo', 'bar'])
1072
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1357
1073
  ```
1074
+ This decorator respects the @project decorator and triggers the flow
1075
+ when upstream runs within the same namespace complete successfully
1358
1076
 
1359
- Additionally, you can specify the parameter mappings
1360
- to map event payload to Metaflow parameters for the flow.
1077
+ Additionally, you can specify project aware upstream flow dependencies
1078
+ by specifying the fully qualified project_flow_name.
1361
1079
  ```
1362
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1080
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1363
1081
  ```
1364
1082
  or
1365
1083
  ```
1366
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1367
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1084
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1368
1085
  ```
1369
1086
 
1370
- 'parameters' can also be a list of strings and tuples like so:
1371
- ```
1372
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1373
- ```
1374
- This is equivalent to:
1087
+ You can also specify just the project or project branch (other values will be
1088
+ inferred from the current project or project branch):
1375
1089
  ```
1376
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1090
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1091
+ ```
1092
+
1093
+ Note that `branch` is typically one of:
1094
+ - `prod`
1095
+ - `user.bob`
1096
+ - `test.my_experiment`
1097
+ - `prod.staging`
1098
+
1099
+
1100
+ Parameters
1101
+ ----------
1102
+ flow : Union[str, Dict[str, str]], optional, default None
1103
+ Upstream flow dependency for this flow.
1104
+ flows : List[Union[str, Dict[str, str]]], default []
1105
+ Upstream flow dependencies for this flow.
1106
+ options : Dict[str, Any], default {}
1107
+ Backend-specific configuration for tuning eventing behavior.
1108
+ """
1109
+ ...
1110
+
1111
+ @typing.overload
1112
+ def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1113
+ ...
1114
+
1115
+ def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}):
1116
+ """
1117
+ Specifies the flow(s) that this flow depends on.
1118
+
1119
+ ```
1120
+ @trigger_on_finish(flow='FooFlow')
1121
+ ```
1122
+ or
1123
+ ```
1124
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1125
+ ```
1126
+ This decorator respects the @project decorator and triggers the flow
1127
+ when upstream runs within the same namespace complete successfully
1128
+
1129
+ Additionally, you can specify project aware upstream flow dependencies
1130
+ by specifying the fully qualified project_flow_name.
1131
+ ```
1132
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1133
+ ```
1134
+ or
1135
+ ```
1136
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1137
+ ```
1138
+
1139
+ You can also specify just the project or project branch (other values will be
1140
+ inferred from the current project or project branch):
1141
+ ```
1142
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1143
+ ```
1144
+
1145
+ Note that `branch` is typically one of:
1146
+ - `prod`
1147
+ - `user.bob`
1148
+ - `test.my_experiment`
1149
+ - `prod.staging`
1150
+
1151
+
1152
+ Parameters
1153
+ ----------
1154
+ flow : Union[str, Dict[str, str]], optional, default None
1155
+ Upstream flow dependency for this flow.
1156
+ flows : List[Union[str, Dict[str, str]]], default []
1157
+ Upstream flow dependencies for this flow.
1158
+ options : Dict[str, Any], default {}
1159
+ Backend-specific configuration for tuning eventing behavior.
1160
+ """
1161
+ ...
1162
+
1163
+ @typing.overload
1164
+ 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]]:
1165
+ """
1166
+ Specifies the event(s) that this flow depends on.
1167
+
1168
+ ```
1169
+ @trigger(event='foo')
1170
+ ```
1171
+ or
1172
+ ```
1173
+ @trigger(events=['foo', 'bar'])
1174
+ ```
1175
+
1176
+ Additionally, you can specify the parameter mappings
1177
+ to map event payload to Metaflow parameters for the flow.
1178
+ ```
1179
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1180
+ ```
1181
+ or
1182
+ ```
1183
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1184
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1185
+ ```
1186
+
1187
+ 'parameters' can also be a list of strings and tuples like so:
1188
+ ```
1189
+ @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1190
+ ```
1191
+ This is equivalent to:
1192
+ ```
1193
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1377
1194
  ```
1378
1195
 
1379
1196
 
@@ -1487,104 +1304,287 @@ def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packa
1487
1304
  """
1488
1305
  ...
1489
1306
 
1490
- @typing.overload
1491
- 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]]:
1307
+ def airflow_external_task_sensor(*, timeout: int, poke_interval: int, mode: str, exponential_backoff: bool, pool: str, soft_fail: bool, name: str, description: str, external_dag_id: str, external_task_ids: typing.List[str], allowed_states: typing.List[str], failed_states: typing.List[str], execution_delta: "datetime.timedelta", check_existence: bool) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1492
1308
  """
1493
- Specifies the flow(s) that this flow depends on.
1309
+ The `@airflow_external_task_sensor` decorator attaches a Airflow [ExternalTaskSensor](https://airflow.apache.org/docs/apache-airflow/stable/_api/airflow/sensors/external_task/index.html#airflow.sensors.external_task.ExternalTaskSensor) before the start step of the flow.
1310
+ This decorator only works when a flow is scheduled on Airflow and is compiled using `airflow create`. More than one `@airflow_external_task_sensor` can be added as a flow decorators. Adding more than one decorator will ensure that `start` step starts only after all sensors finish.
1494
1311
 
1495
- ```
1496
- @trigger_on_finish(flow='FooFlow')
1497
- ```
1498
- or
1499
- ```
1500
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1501
- ```
1502
- This decorator respects the @project decorator and triggers the flow
1503
- when upstream runs within the same namespace complete successfully
1504
1312
 
1505
- Additionally, you can specify project aware upstream flow dependencies
1506
- by specifying the fully qualified project_flow_name.
1507
- ```
1508
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1509
- ```
1510
- or
1511
- ```
1512
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1513
- ```
1313
+ Parameters
1314
+ ----------
1315
+ timeout : int
1316
+ Time, in seconds before the task times out and fails. (Default: 3600)
1317
+ poke_interval : int
1318
+ Time in seconds that the job should wait in between each try. (Default: 60)
1319
+ mode : str
1320
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1321
+ exponential_backoff : bool
1322
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1323
+ pool : str
1324
+ the slot pool this task should run in,
1325
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1326
+ soft_fail : bool
1327
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1328
+ name : str
1329
+ Name of the sensor on Airflow
1330
+ description : str
1331
+ Description of sensor in the Airflow UI
1332
+ external_dag_id : str
1333
+ The dag_id that contains the task you want to wait for.
1334
+ external_task_ids : List[str]
1335
+ The list of task_ids that you want to wait for.
1336
+ If None (default value) the sensor waits for the DAG. (Default: None)
1337
+ allowed_states : List[str]
1338
+ Iterable of allowed states, (Default: ['success'])
1339
+ failed_states : List[str]
1340
+ Iterable of failed or dis-allowed states. (Default: None)
1341
+ execution_delta : datetime.timedelta
1342
+ time difference with the previous execution to look at,
1343
+ the default is the same logical date as the current task or DAG. (Default: None)
1344
+ check_existence: bool
1345
+ Set to True to check if the external task exists or check if
1346
+ the DAG to wait for exists. (Default: True)
1347
+ """
1348
+ ...
1349
+
1350
+ def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None):
1351
+ """
1352
+ Allows setting external datastores to save data for the
1353
+ `@checkpoint`/`@model`/`@huggingface_hub` decorators.
1514
1354
 
1515
- You can also specify just the project or project branch (other values will be
1516
- inferred from the current project or project branch):
1517
- ```
1518
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1519
- ```
1355
+ This decorator is useful when users wish to save data to a different datastore
1356
+ than what is configured in Metaflow. This can be for variety of reasons:
1520
1357
 
1521
- Note that `branch` is typically one of:
1522
- - `prod`
1523
- - `user.bob`
1524
- - `test.my_experiment`
1525
- - `prod.staging`
1358
+ 1. Data security: The objects needs to be stored in a bucket (object storage) that is not accessible by other flows.
1359
+ 2. Data Locality: The location where the task is executing is not located in the same region as the datastore.
1360
+ - Example: Metaflow datastore lives in US East, but the task is executing in Finland datacenters.
1361
+ 3. Data Lifecycle Policies: The objects need to be archived / managed separately from the Metaflow managed objects.
1362
+ - Example: Flow is training very large models that need to be stored separately and will be deleted more aggressively than the Metaflow managed objects.
1363
+
1364
+ Usage:
1365
+ ----------
1366
+
1367
+ - Using a custom IAM role to access the datastore.
1368
+
1369
+ ```python
1370
+ @with_artifact_store(
1371
+ type="s3",
1372
+ config=lambda: {
1373
+ "root": "s3://my-bucket-foo/path/to/root",
1374
+ "role_arn": ROLE,
1375
+ },
1376
+ )
1377
+ class MyFlow(FlowSpec):
1378
+
1379
+ @checkpoint
1380
+ @step
1381
+ def start(self):
1382
+ with open("my_file.txt", "w") as f:
1383
+ f.write("Hello, World!")
1384
+ self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1385
+ self.next(self.end)
1386
+
1387
+ ```
1388
+
1389
+ - Using credentials to access the s3-compatible datastore.
1390
+
1391
+ ```python
1392
+ @with_artifact_store(
1393
+ type="s3",
1394
+ config=lambda: {
1395
+ "root": "s3://my-bucket-foo/path/to/root",
1396
+ "client_params": {
1397
+ "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1398
+ "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1399
+ },
1400
+ },
1401
+ )
1402
+ class MyFlow(FlowSpec):
1403
+
1404
+ @checkpoint
1405
+ @step
1406
+ def start(self):
1407
+ with open("my_file.txt", "w") as f:
1408
+ f.write("Hello, World!")
1409
+ self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1410
+ self.next(self.end)
1411
+
1412
+ ```
1413
+
1414
+ - Accessing objects stored in external datastores after task execution.
1415
+
1416
+ ```python
1417
+ run = Run("CheckpointsTestsFlow/8992")
1418
+ with artifact_store_from(run=run, config={
1419
+ "client_params": {
1420
+ "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1421
+ "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1422
+ },
1423
+ }):
1424
+ with Checkpoint() as cp:
1425
+ latest = cp.list(
1426
+ task=run["start"].task
1427
+ )[0]
1428
+ print(latest)
1429
+ cp.load(
1430
+ latest,
1431
+ "test-checkpoints"
1432
+ )
1433
+
1434
+ task = Task("TorchTuneFlow/8484/train/53673")
1435
+ with artifact_store_from(run=run, config={
1436
+ "client_params": {
1437
+ "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1438
+ "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1439
+ },
1440
+ }):
1441
+ load_model(
1442
+ task.data.model_ref,
1443
+ "test-models"
1444
+ )
1445
+ ```
1446
+ Parameters:
1447
+ ----------
1448
+
1449
+ type: str
1450
+ The type of the datastore. Can be one of 's3', 'gcs', 'azure' or any other supported metaflow Datastore.
1451
+
1452
+ config: dict or Callable
1453
+ Dictionary of configuration options for the datastore. The following keys are required:
1454
+ - root: The root path in the datastore where the data will be saved. (needs to be in the format expected by the datastore)
1455
+ - example: 's3://bucket-name/path/to/root'
1456
+ - example: 'gs://bucket-name/path/to/root'
1457
+ - example: 'https://myblockacc.blob.core.windows.net/metaflow/'
1458
+ - role_arn (optional): AWS IAM role to access s3 bucket (only when `type` is 's3')
1459
+ - session_vars (optional): AWS session variables to access s3 bucket (only when `type` is 's3')
1460
+ - client_params (optional): AWS client parameters to access s3 bucket (only when `type` is 's3')
1461
+ """
1462
+ ...
1463
+
1464
+ def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1465
+ """
1466
+ Specifies what flows belong to the same project.
1467
+
1468
+ A project-specific namespace is created for all flows that
1469
+ use the same `@project(name)`.
1526
1470
 
1527
1471
 
1528
1472
  Parameters
1529
1473
  ----------
1530
- flow : Union[str, Dict[str, str]], optional, default None
1531
- Upstream flow dependency for this flow.
1532
- flows : List[Union[str, Dict[str, str]]], default []
1533
- Upstream flow dependencies for this flow.
1534
- options : Dict[str, Any], default {}
1535
- Backend-specific configuration for tuning eventing behavior.
1474
+ name : str
1475
+ Project name. Make sure that the name is unique amongst all
1476
+ projects that use the same production scheduler. The name may
1477
+ contain only lowercase alphanumeric characters and underscores.
1478
+
1479
+ branch : Optional[str], default None
1480
+ The branch to use. If not specified, the branch is set to
1481
+ `user.<username>` unless `production` is set to `True`. This can
1482
+ also be set on the command line using `--branch` as a top-level option.
1483
+ It is an error to specify `branch` in the decorator and on the command line.
1484
+
1485
+ production : bool, default False
1486
+ Whether or not the branch is the production branch. This can also be set on the
1487
+ command line using `--production` as a top-level option. It is an error to specify
1488
+ `production` in the decorator and on the command line.
1489
+ The project branch name will be:
1490
+ - if `branch` is specified:
1491
+ - if `production` is True: `prod.<branch>`
1492
+ - if `production` is False: `test.<branch>`
1493
+ - if `branch` is not specified:
1494
+ - if `production` is True: `prod`
1495
+ - if `production` is False: `user.<username>`
1536
1496
  """
1537
1497
  ...
1538
1498
 
1539
1499
  @typing.overload
1540
- def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1500
+ def schedule(*, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1501
+ """
1502
+ Specifies the times when the flow should be run when running on a
1503
+ production scheduler.
1504
+
1505
+
1506
+ Parameters
1507
+ ----------
1508
+ hourly : bool, default False
1509
+ Run the workflow hourly.
1510
+ daily : bool, default True
1511
+ Run the workflow daily.
1512
+ weekly : bool, default False
1513
+ Run the workflow weekly.
1514
+ cron : str, optional, default None
1515
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1516
+ specified by this expression.
1517
+ timezone : str, optional, default None
1518
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1519
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1520
+ """
1541
1521
  ...
1542
1522
 
1543
- def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}):
1523
+ @typing.overload
1524
+ def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1525
+ ...
1526
+
1527
+ def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None):
1544
1528
  """
1545
- Specifies the flow(s) that this flow depends on.
1529
+ Specifies the times when the flow should be run when running on a
1530
+ production scheduler.
1546
1531
 
1547
- ```
1548
- @trigger_on_finish(flow='FooFlow')
1549
- ```
1550
- or
1551
- ```
1552
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1553
- ```
1554
- This decorator respects the @project decorator and triggers the flow
1555
- when upstream runs within the same namespace complete successfully
1556
1532
 
1557
- Additionally, you can specify project aware upstream flow dependencies
1558
- by specifying the fully qualified project_flow_name.
1559
- ```
1560
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1561
- ```
1562
- or
1563
- ```
1564
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1565
- ```
1533
+ Parameters
1534
+ ----------
1535
+ hourly : bool, default False
1536
+ Run the workflow hourly.
1537
+ daily : bool, default True
1538
+ Run the workflow daily.
1539
+ weekly : bool, default False
1540
+ Run the workflow weekly.
1541
+ cron : str, optional, default None
1542
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1543
+ specified by this expression.
1544
+ timezone : str, optional, default None
1545
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1546
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1547
+ """
1548
+ ...
1549
+
1550
+ @typing.overload
1551
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1552
+ """
1553
+ Specifies the PyPI packages for all steps of the flow.
1566
1554
 
1567
- You can also specify just the project or project branch (other values will be
1568
- inferred from the current project or project branch):
1569
- ```
1570
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1571
- ```
1555
+ Use `@pypi_base` to set common packages required by all
1556
+ steps and use `@pypi` to specify step-specific overrides.
1572
1557
 
1573
- Note that `branch` is typically one of:
1574
- - `prod`
1575
- - `user.bob`
1576
- - `test.my_experiment`
1577
- - `prod.staging`
1558
+ Parameters
1559
+ ----------
1560
+ packages : Dict[str, str], default: {}
1561
+ Packages to use for this flow. The key is the name of the package
1562
+ and the value is the version to use.
1563
+ python : str, optional, default: None
1564
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1565
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1566
+ """
1567
+ ...
1568
+
1569
+ @typing.overload
1570
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1571
+ ...
1572
+
1573
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1574
+ """
1575
+ Specifies the PyPI packages for all steps of the flow.
1578
1576
 
1577
+ Use `@pypi_base` to set common packages required by all
1578
+ steps and use `@pypi` to specify step-specific overrides.
1579
1579
 
1580
1580
  Parameters
1581
1581
  ----------
1582
- flow : Union[str, Dict[str, str]], optional, default None
1583
- Upstream flow dependency for this flow.
1584
- flows : List[Union[str, Dict[str, str]]], default []
1585
- Upstream flow dependencies for this flow.
1586
- options : Dict[str, Any], default {}
1587
- Backend-specific configuration for tuning eventing behavior.
1582
+ packages : Dict[str, str], default: {}
1583
+ Packages to use for this flow. The key is the name of the package
1584
+ and the value is the version to use.
1585
+ python : str, optional, default: None
1586
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1587
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1588
1588
  """
1589
1589
  ...
1590
1590