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