ob-metaflow-stubs 6.0.3.180rc4__py2.py3-none-any.whl → 6.0.3.181__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 +739 -740
  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 +5 -5
  8. metaflow-stubs/client/filecache.pyi +2 -2
  9. metaflow-stubs/events.pyi +2 -2
  10. metaflow-stubs/exception.pyi +1 -1
  11. metaflow-stubs/flowspec.pyi +4 -4
  12. metaflow-stubs/generated_for.txt +1 -1
  13. metaflow-stubs/includefile.pyi +2 -2
  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 +2 -2
  18. metaflow-stubs/metadata_provider/util.pyi +1 -1
  19. metaflow-stubs/metaflow_config.pyi +1 -1
  20. metaflow-stubs/metaflow_current.pyi +80 -80
  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 +1 -1
  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 +3 -3
  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 +3 -3
  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 +4 -4
  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 +2 -2
  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 +2 -2
  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 +1 -1
  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 +3 -3
  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 +2 -2
  63. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/tar_utils.pyi +2 -2
  64. metaflow-stubs/mf_extensions/outerbounds/__init__.pyi +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 +1 -1
  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 +12 -2
  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 +13 -13
  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 +3 -3
  102. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +2 -2
  103. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +3 -3
  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 +2 -2
  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 +3 -3
  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 +2 -2
  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 +2 -2
  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 +4 -4
  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 +2 -2
  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 +2 -2
  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 +2 -2
  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 +4 -4
  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 +1 -1
  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 +2 -2
  188. metaflow-stubs/plugins/timeout_decorator.pyi +2 -2
  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 +29 -29
  196. metaflow-stubs/runner/deployer_impl.pyi +2 -2
  197. metaflow-stubs/runner/metaflow_runner.pyi +3 -3
  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 +3 -3
  202. metaflow-stubs/system/__init__.pyi +1 -1
  203. metaflow-stubs/system/system_logger.pyi +1 -1
  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 +5 -5
  209. metaflow-stubs/user_configs/config_options.pyi +3 -3
  210. metaflow-stubs/user_configs/config_parameters.pyi +4 -4
  211. {ob_metaflow_stubs-6.0.3.180rc4.dist-info → ob_metaflow_stubs-6.0.3.181.dist-info}/METADATA +1 -1
  212. ob_metaflow_stubs-6.0.3.181.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.180rc4.dist-info/RECORD +0 -220
  219. {ob_metaflow_stubs-6.0.3.180rc4.dist-info → ob_metaflow_stubs-6.0.3.181.dist-info}/WHEEL +0 -0
  220. {ob_metaflow_stubs-6.0.3.180rc4.dist-info → ob_metaflow_stubs-6.0.3.181.dist-info}/top_level.txt +0 -0
@@ -1,7 +1,7 @@
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-17T18:51:01.431227 #
4
+ # Generated on 2025-06-18T08:04:08.962280 #
5
5
  ######################################################################################################
6
6
 
7
7
  from __future__ import annotations
@@ -36,9 +36,9 @@ 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
- from . import tuple_util as tuple_util
41
39
  from . import events as events
40
+ from . import tuple_util as tuple_util
41
+ from . import metaflow_git as metaflow_git
42
42
  from . import runner as runner
43
43
  from . import plugins as plugins
44
44
  from .mf_extensions.outerbounds.toplevel.global_aliases_for_metaflow_package import S3 as S3
@@ -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,127 +154,54 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
155
154
  """
156
155
  ...
157
156
 
158
- 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]]]:
159
- """
160
- This decorator is used to run Ollama APIs as Metaflow task sidecars.
161
-
162
- User code call
163
- --------------
164
- @ollama(
165
- models=[...],
166
- ...
167
- )
168
-
169
- Valid backend options
170
- ---------------------
171
- - 'local': Run as a separate process on the local task machine.
172
- - (TODO) 'managed': Outerbounds hosts and selects compute provider.
173
- - (TODO) 'remote': Spin up separate instance to serve Ollama models.
174
-
175
- Valid model options
176
- -------------------
177
- Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
178
-
179
-
180
- Parameters
181
- ----------
182
- models: list[str]
183
- List of Ollama containers running models in sidecars.
184
- backend: str
185
- Determines where and how to run the Ollama process.
186
- force_pull: bool
187
- Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
188
- cache_update_policy: str
189
- Cache update policy: "auto", "force", or "never".
190
- force_cache_update: bool
191
- Simple override for "force" cache update policy.
192
- debug: bool
193
- Whether to turn on verbose debugging logs.
194
- circuit_breaker_config: dict
195
- Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
196
- timeout_config: dict
197
- Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
198
- """
199
- ...
200
-
201
157
  @typing.overload
202
- 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]]]:
158
+ 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]]]:
203
159
  """
204
- Specifies the PyPI packages for the step.
160
+ Specifies that the step will success under all circumstances.
205
161
 
206
- Information in this decorator will augment any
207
- attributes set in the `@pyi_base` flow-level decorator. Hence,
208
- you can use `@pypi_base` to set packages required by all
209
- steps and use `@pypi` to specify step-specific overrides.
162
+ The decorator will create an optional artifact, specified by `var`, which
163
+ contains the exception raised. You can use it to detect the presence
164
+ of errors, indicating that all happy-path artifacts produced by the step
165
+ are missing.
210
166
 
211
167
 
212
168
  Parameters
213
169
  ----------
214
- packages : Dict[str, str], default: {}
215
- Packages to use for this step. The key is the name of the package
216
- and the value is the version to use.
217
- python : str, optional, default: None
218
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
219
- that the version used will correspond to the version of the Python interpreter used to start the run.
170
+ var : str, optional, default None
171
+ Name of the artifact in which to store the caught exception.
172
+ If not specified, the exception is not stored.
173
+ print_exception : bool, default True
174
+ Determines whether or not the exception is printed to
175
+ stdout when caught.
220
176
  """
221
177
  ...
222
178
 
223
179
  @typing.overload
224
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
180
+ def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
225
181
  ...
226
182
 
227
183
  @typing.overload
228
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
229
- ...
230
-
231
- 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):
232
- """
233
- Specifies the PyPI packages for the step.
234
-
235
- Information in this decorator will augment any
236
- attributes set in the `@pyi_base` flow-level decorator. Hence,
237
- you can use `@pypi_base` to set packages required by all
238
- steps and use `@pypi` to specify step-specific overrides.
239
-
240
-
241
- Parameters
242
- ----------
243
- packages : Dict[str, str], default: {}
244
- Packages to use for this step. The key is the name of the package
245
- and the value is the version to use.
246
- python : str, optional, default: None
247
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
248
- that the version used will correspond to the version of the Python interpreter used to start the run.
249
- """
184
+ def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
250
185
  ...
251
186
 
252
- def nvidia(*, gpu: int, gpu_type: str, queue_timeout: int) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
187
+ 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):
253
188
  """
254
- Specifies that this step should execute on DGX cloud.
255
-
189
+ Specifies that the step will success under all circumstances.
256
190
 
257
- Parameters
258
- ----------
259
- gpu : int
260
- Number of GPUs to use.
261
- gpu_type : str
262
- Type of Nvidia GPU to use.
263
- queue_timeout : int
264
- Time to keep the job in NVCF's queue.
265
- """
266
- ...
267
-
268
- 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]]]:
269
- """
270
- Specifies that this step should execute on DGX cloud.
191
+ The decorator will create an optional artifact, specified by `var`, which
192
+ contains the exception raised. You can use it to detect the presence
193
+ of errors, indicating that all happy-path artifacts produced by the step
194
+ are missing.
271
195
 
272
196
 
273
197
  Parameters
274
198
  ----------
275
- gpu : int
276
- Number of GPUs to use.
277
- gpu_type : str
278
- Type of Nvidia GPU to use.
199
+ var : str, optional, default None
200
+ Name of the artifact in which to store the caught exception.
201
+ If not specified, the exception is not stored.
202
+ print_exception : bool, default True
203
+ Determines whether or not the exception is printed to
204
+ stdout when caught.
279
205
  """
280
206
  ...
281
207
 
@@ -335,125 +261,237 @@ def retry(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
335
261
  ...
336
262
 
337
263
  @typing.overload
338
- 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]]]:
264
+ 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]]]:
339
265
  """
340
- Specifies that the step will success under all circumstances.
341
-
342
- The decorator will create an optional artifact, specified by `var`, which
343
- contains the exception raised. You can use it to detect the presence
344
- of errors, indicating that all happy-path artifacts produced by the step
345
- are missing.
266
+ Specifies environment variables to be set prior to the execution of a step.
346
267
 
347
268
 
348
269
  Parameters
349
270
  ----------
350
- var : str, optional, default None
351
- Name of the artifact in which to store the caught exception.
352
- If not specified, the exception is not stored.
353
- print_exception : bool, default True
354
- Determines whether or not the exception is printed to
355
- stdout when caught.
271
+ vars : Dict[str, str], default {}
272
+ Dictionary of environment variables to set.
356
273
  """
357
274
  ...
358
275
 
359
276
  @typing.overload
360
- def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
277
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
361
278
  ...
362
279
 
363
280
  @typing.overload
364
- def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
281
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
365
282
  ...
366
283
 
367
- 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):
284
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
368
285
  """
369
- Specifies that the step will success under all circumstances.
370
-
371
- The decorator will create an optional artifact, specified by `var`, which
372
- contains the exception raised. You can use it to detect the presence
373
- of errors, indicating that all happy-path artifacts produced by the step
374
- are missing.
286
+ Specifies environment variables to be set prior to the execution of a step.
375
287
 
376
288
 
377
289
  Parameters
378
290
  ----------
379
- var : str, optional, default None
380
- Name of the artifact in which to store the caught exception.
381
- If not specified, the exception is not stored.
382
- print_exception : bool, default True
383
- Determines whether or not the exception is printed to
384
- stdout when caught.
291
+ vars : Dict[str, str], default {}
292
+ Dictionary of environment variables to set.
385
293
  """
386
294
  ...
387
295
 
388
- @typing.overload
389
- 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]]]:
296
+ 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]]]:
390
297
  """
391
- Specifies secrets to be retrieved and injected as environment variables prior to
392
- the execution of a step.
298
+ This decorator is used to run vllm APIs as Metaflow task sidecars.
299
+
300
+ User code call
301
+ --------------
302
+ @vllm(
303
+ model="...",
304
+ ...
305
+ )
306
+
307
+ Valid backend options
308
+ ---------------------
309
+ - 'local': Run as a separate process on the local task machine.
310
+
311
+ Valid model options
312
+ -------------------
313
+ Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
314
+
315
+ NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
316
+ If you need multiple models, you must create multiple @vllm decorators.
393
317
 
394
318
 
395
319
  Parameters
396
320
  ----------
397
- sources : List[Union[str, Dict[str, Any]]], default: []
398
- List of secret specs, defining how the secrets are to be retrieved
321
+ model: str
322
+ HuggingFace model identifier to be served by vLLM.
323
+ backend: str
324
+ Determines where and how to run the vLLM process.
325
+ debug: bool
326
+ Whether to turn on verbose debugging logs.
327
+ kwargs : Any
328
+ Any other keyword arguments are passed directly to the vLLM engine.
329
+ This allows for flexible configuration of vLLM server settings.
330
+ For example, `tensor_parallel_size=2`.
399
331
  """
400
332
  ...
401
333
 
402
- @typing.overload
403
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
404
- ...
405
-
406
- @typing.overload
407
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
408
- ...
409
-
410
- def secrets(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = []):
334
+ 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]]]:
411
335
  """
412
- Specifies secrets to be retrieved and injected as environment variables prior to
413
- the execution of a step.
336
+ Decorator that helps cache, version and store models/datasets from huggingface hub.
414
337
 
415
338
 
416
339
  Parameters
417
340
  ----------
418
- sources : List[Union[str, Dict[str, Any]]], default: []
419
- List of secret specs, defining how the secrets are to be retrieved
341
+ temp_dir_root : str, optional
342
+ The root directory that will hold the temporary directory where objects will be downloaded.
343
+
344
+ load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
345
+ The list of repos (models/datasets) to load.
346
+
347
+ Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
348
+
349
+ - If repo (model/dataset) is not found in the datastore:
350
+ - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
351
+ - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
352
+ - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
353
+
354
+ - If repo is found in the datastore:
355
+ - Loads it directly from datastore to local path (can be temporary directory or specified path)
420
356
  """
421
357
  ...
422
358
 
423
359
  @typing.overload
424
- 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]]]:
360
+ def checkpoint(*, load_policy: str = 'fresh', temp_dir_root: str = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
425
361
  """
426
- Enables loading / saving of models within a step.
362
+ Enables checkpointing for a step.
427
363
 
428
364
 
429
365
 
430
366
  Parameters
431
367
  ----------
432
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
433
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
434
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
435
- - `current.checkpoint`
436
- - `current.model`
437
- - `current.huggingface_hub`
368
+ load_policy : str, default: "fresh"
369
+ The policy for loading the checkpoint. The following policies are supported:
370
+ - "eager": Loads the the latest available checkpoint within the namespace.
371
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
372
+ will be loaded at the start of the task.
373
+ - "none": Do not load any checkpoint
374
+ - "fresh": Loads the lastest checkpoint created within the running Task.
375
+ This mode helps loading checkpoints across various retry attempts of the same task.
376
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
377
+ created within the task will be loaded when the task is retries execution on failure.
438
378
 
439
- 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
440
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
441
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
379
+ temp_dir_root : str, default: None
380
+ The root directory under which `current.checkpoint.directory` will be created.
381
+ """
382
+ ...
383
+
384
+ @typing.overload
385
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
386
+ ...
387
+
388
+ @typing.overload
389
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
390
+ ...
391
+
392
+ def checkpoint(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, load_policy: str = 'fresh', temp_dir_root: str = None):
393
+ """
394
+ Enables checkpointing for a step.
395
+
396
+
397
+
398
+ Parameters
399
+ ----------
400
+ load_policy : str, default: "fresh"
401
+ The policy for loading the checkpoint. The following policies are supported:
402
+ - "eager": Loads the the latest available checkpoint within the namespace.
403
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
404
+ will be loaded at the start of the task.
405
+ - "none": Do not load any checkpoint
406
+ - "fresh": Loads the lastest checkpoint created within the running Task.
407
+ This mode helps loading checkpoints across various retry attempts of the same task.
408
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
409
+ created within the task will be loaded when the task is retries execution on failure.
442
410
 
443
411
  temp_dir_root : str, default: None
444
- The root directory under which `current.model.loaded` will store loaded models
412
+ The root directory under which `current.checkpoint.directory` will be created.
413
+ """
414
+ ...
415
+
416
+ 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]]]:
417
+ """
418
+ This decorator is used to run Ollama APIs as Metaflow task sidecars.
419
+
420
+ User code call
421
+ --------------
422
+ @ollama(
423
+ models=[...],
424
+ ...
425
+ )
426
+
427
+ Valid backend options
428
+ ---------------------
429
+ - 'local': Run as a separate process on the local task machine.
430
+ - (TODO) 'managed': Outerbounds hosts and selects compute provider.
431
+ - (TODO) 'remote': Spin up separate instance to serve Ollama models.
432
+
433
+ Valid model options
434
+ -------------------
435
+ Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
436
+
437
+
438
+ Parameters
439
+ ----------
440
+ models: list[str]
441
+ List of Ollama containers running models in sidecars.
442
+ backend: str
443
+ Determines where and how to run the Ollama process.
444
+ force_pull: bool
445
+ Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
446
+ cache_update_policy: str
447
+ Cache update policy: "auto", "force", or "never".
448
+ force_cache_update: bool
449
+ Simple override for "force" cache update policy.
450
+ debug: bool
451
+ Whether to turn on verbose debugging logs.
452
+ circuit_breaker_config: dict
453
+ Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
454
+ timeout_config: dict
455
+ Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
456
+ """
457
+ ...
458
+
459
+ 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]]]:
460
+ """
461
+ Specifies that this step is used to deploy an instance of the app.
462
+ Requires that self.app_name, self.app_port, self.entrypoint and self.deployDir is set.
463
+
464
+
465
+ Parameters
466
+ ----------
467
+ app_port : int
468
+ Number of GPUs to use.
469
+ app_name : str
470
+ Name of the app to deploy.
445
471
  """
446
472
  ...
447
473
 
448
474
  @typing.overload
449
- def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
475
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
476
+ """
477
+ Decorator prototype for all step decorators. This function gets specialized
478
+ and imported for all decorators types by _import_plugin_decorators().
479
+ """
450
480
  ...
451
481
 
452
482
  @typing.overload
453
- def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
483
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
454
484
  ...
455
485
 
456
- 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):
486
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
487
+ """
488
+ Decorator prototype for all step decorators. This function gets specialized
489
+ and imported for all decorators types by _import_plugin_decorators().
490
+ """
491
+ ...
492
+
493
+ @typing.overload
494
+ 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]]]:
457
495
  """
458
496
  Enables loading / saving of models within a step.
459
497
 
@@ -477,346 +515,121 @@ def model(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
477
515
  """
478
516
  ...
479
517
 
480
- 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]]]:
518
+ @typing.overload
519
+ def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
520
+ ...
521
+
522
+ @typing.overload
523
+ def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
524
+ ...
525
+
526
+ 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):
481
527
  """
482
- Decorator that helps cache, version and store models/datasets from huggingface hub.
528
+ Enables loading / saving of models within a step.
529
+
483
530
 
484
531
 
485
532
  Parameters
486
533
  ----------
487
- temp_dir_root : str, optional
488
- The root directory that will hold the temporary directory where objects will be downloaded.
489
-
490
- load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
491
- The list of repos (models/datasets) to load.
492
-
493
- Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
534
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
535
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
536
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
537
+ - `current.checkpoint`
538
+ - `current.model`
539
+ - `current.huggingface_hub`
494
540
 
495
- - If repo (model/dataset) is not found in the datastore:
496
- - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
497
- - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
498
- - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
541
+ 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
542
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
543
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
499
544
 
500
- - If repo is found in the datastore:
501
- - Loads it directly from datastore to local path (can be temporary directory or specified path)
545
+ temp_dir_root : str, default: None
546
+ The root directory under which `current.model.loaded` will store loaded models
502
547
  """
503
548
  ...
504
549
 
505
550
  @typing.overload
506
- 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]]]:
551
+ 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]]]:
507
552
  """
508
- Specifies the resources needed when executing this step.
509
-
510
- Use `@resources` to specify the resource requirements
511
- independently of the specific compute layer (`@batch`, `@kubernetes`).
512
-
513
- You can choose the compute layer on the command line by executing e.g.
514
- ```
515
- python myflow.py run --with batch
516
- ```
517
- or
518
- ```
519
- python myflow.py run --with kubernetes
520
- ```
521
- which executes the flow on the desired system using the
522
- requirements specified in `@resources`.
553
+ Specifies secrets to be retrieved and injected as environment variables prior to
554
+ the execution of a step.
523
555
 
524
556
 
525
557
  Parameters
526
558
  ----------
527
- cpu : int, default 1
528
- Number of CPUs required for this step.
529
- gpu : int, optional, default None
530
- Number of GPUs required for this step.
531
- disk : int, optional, default None
532
- Disk size (in MB) required for this step. Only applies on Kubernetes.
533
- memory : int, default 4096
534
- Memory size (in MB) required for this step.
535
- shared_memory : int, optional, default None
536
- The value for the size (in MiB) of the /dev/shm volume for this step.
537
- This parameter maps to the `--shm-size` option in Docker.
559
+ sources : List[Union[str, Dict[str, Any]]], default: []
560
+ List of secret specs, defining how the secrets are to be retrieved
538
561
  """
539
562
  ...
540
563
 
541
564
  @typing.overload
542
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
565
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
543
566
  ...
544
567
 
545
568
  @typing.overload
546
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
569
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
547
570
  ...
548
571
 
549
- def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None):
572
+ def secrets(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = []):
550
573
  """
551
- Specifies the resources needed when executing this step.
552
-
553
- Use `@resources` to specify the resource requirements
554
- independently of the specific compute layer (`@batch`, `@kubernetes`).
555
-
556
- You can choose the compute layer on the command line by executing e.g.
557
- ```
558
- python myflow.py run --with batch
559
- ```
560
- or
561
- ```
562
- python myflow.py run --with kubernetes
563
- ```
564
- which executes the flow on the desired system using the
565
- requirements specified in `@resources`.
574
+ Specifies secrets to be retrieved and injected as environment variables prior to
575
+ the execution of a step.
566
576
 
567
577
 
568
578
  Parameters
569
579
  ----------
570
- cpu : int, default 1
571
- Number of CPUs required for this step.
572
- gpu : int, optional, default None
573
- Number of GPUs required for this step.
574
- disk : int, optional, default None
575
- Disk size (in MB) required for this step. Only applies on Kubernetes.
576
- memory : int, default 4096
577
- Memory size (in MB) required for this step.
578
- shared_memory : int, optional, default None
579
- The value for the size (in MiB) of the /dev/shm volume for this step.
580
- This parameter maps to the `--shm-size` option in Docker.
580
+ sources : List[Union[str, Dict[str, Any]]], default: []
581
+ List of secret specs, defining how the secrets are to be retrieved
581
582
  """
582
583
  ...
583
584
 
584
585
  @typing.overload
585
- 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]]]:
586
+ 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]]]:
586
587
  """
587
- Specifies the Conda environment for the step.
588
+ Specifies the PyPI packages for the step.
588
589
 
589
590
  Information in this decorator will augment any
590
- attributes set in the `@conda_base` flow-level decorator. Hence,
591
- you can use `@conda_base` to set packages required by all
592
- steps and use `@conda` to specify step-specific overrides.
591
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
592
+ you can use `@pypi_base` to set packages required by all
593
+ steps and use `@pypi` to specify step-specific overrides.
593
594
 
594
595
 
595
596
  Parameters
596
597
  ----------
597
- packages : Dict[str, str], default {}
598
+ packages : Dict[str, str], default: {}
598
599
  Packages to use for this step. The key is the name of the package
599
600
  and the value is the version to use.
600
- libraries : Dict[str, str], default {}
601
- Supported for backward compatibility. When used with packages, packages will take precedence.
602
- python : str, optional, default None
601
+ python : str, optional, default: None
603
602
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
604
603
  that the version used will correspond to the version of the Python interpreter used to start the run.
605
- disabled : bool, default False
606
- If set to True, disables @conda.
607
604
  """
608
605
  ...
609
606
 
610
607
  @typing.overload
611
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
608
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
612
609
  ...
613
610
 
614
611
  @typing.overload
615
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
612
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
616
613
  ...
617
614
 
618
- def conda(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
615
+ 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):
619
616
  """
620
- Specifies the Conda environment for the step.
617
+ Specifies the PyPI packages for the step.
621
618
 
622
619
  Information in this decorator will augment any
623
- attributes set in the `@conda_base` flow-level decorator. Hence,
624
- you can use `@conda_base` to set packages required by all
625
- steps and use `@conda` to specify step-specific overrides.
620
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
621
+ you can use `@pypi_base` to set packages required by all
622
+ steps and use `@pypi` to specify step-specific overrides.
626
623
 
627
624
 
628
625
  Parameters
629
626
  ----------
630
- packages : Dict[str, str], default {}
627
+ packages : Dict[str, str], default: {}
631
628
  Packages to use for this step. The key is the name of the package
632
629
  and the value is the version to use.
633
- libraries : Dict[str, str], default {}
634
- Supported for backward compatibility. When used with packages, packages will take precedence.
635
- python : str, optional, default None
630
+ python : str, optional, default: None
636
631
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
637
632
  that the version used will correspond to the version of the Python interpreter used to start the run.
638
- disabled : bool, default False
639
- If set to True, disables @conda.
640
- """
641
- ...
642
-
643
- @typing.overload
644
- def checkpoint(*, load_policy: str = 'fresh', temp_dir_root: str = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
645
- """
646
- Enables checkpointing for a step.
647
-
648
-
649
-
650
- Parameters
651
- ----------
652
- load_policy : str, default: "fresh"
653
- The policy for loading the checkpoint. The following policies are supported:
654
- - "eager": Loads the the latest available checkpoint within the namespace.
655
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
656
- will be loaded at the start of the task.
657
- - "none": Do not load any checkpoint
658
- - "fresh": Loads the lastest checkpoint created within the running Task.
659
- This mode helps loading checkpoints across various retry attempts of the same task.
660
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
661
- created within the task will be loaded when the task is retries execution on failure.
662
-
663
- temp_dir_root : str, default: None
664
- The root directory under which `current.checkpoint.directory` will be created.
665
- """
666
- ...
667
-
668
- @typing.overload
669
- def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
670
- ...
671
-
672
- @typing.overload
673
- def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
674
- ...
675
-
676
- def checkpoint(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, load_policy: str = 'fresh', temp_dir_root: str = None):
677
- """
678
- Enables checkpointing for a step.
679
-
680
-
681
-
682
- Parameters
683
- ----------
684
- load_policy : str, default: "fresh"
685
- The policy for loading the checkpoint. The following policies are supported:
686
- - "eager": Loads the the latest available checkpoint within the namespace.
687
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
688
- will be loaded at the start of the task.
689
- - "none": Do not load any checkpoint
690
- - "fresh": Loads the lastest checkpoint created within the running Task.
691
- This mode helps loading checkpoints across various retry attempts of the same task.
692
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
693
- created within the task will be loaded when the task is retries execution on failure.
694
-
695
- temp_dir_root : str, default: None
696
- The root directory under which `current.checkpoint.directory` will be created.
697
- """
698
- ...
699
-
700
- @typing.overload
701
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
702
- """
703
- Decorator prototype for all step decorators. This function gets specialized
704
- and imported for all decorators types by _import_plugin_decorators().
705
- """
706
- ...
707
-
708
- @typing.overload
709
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
710
- ...
711
-
712
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
713
- """
714
- Decorator prototype for all step decorators. This function gets specialized
715
- and imported for all decorators types by _import_plugin_decorators().
716
- """
717
- ...
718
-
719
- 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]]]:
720
- """
721
- This decorator is used to run vllm APIs as Metaflow task sidecars.
722
-
723
- User code call
724
- --------------
725
- @vllm(
726
- model="...",
727
- ...
728
- )
729
-
730
- Valid backend options
731
- ---------------------
732
- - 'local': Run as a separate process on the local task machine.
733
-
734
- Valid model options
735
- -------------------
736
- Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
737
-
738
- NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
739
- If you need multiple models, you must create multiple @vllm decorators.
740
-
741
-
742
- Parameters
743
- ----------
744
- model: str
745
- HuggingFace model identifier to be served by vLLM.
746
- backend: str
747
- Determines where and how to run the vLLM process.
748
- debug: bool
749
- Whether to turn on verbose debugging logs.
750
- kwargs : Any
751
- Any other keyword arguments are passed directly to the vLLM engine.
752
- This allows for flexible configuration of vLLM server settings.
753
- For example, `tensor_parallel_size=2`.
754
- """
755
- ...
756
-
757
- @typing.overload
758
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
759
- """
760
- Internal decorator to support Fast bakery
761
- """
762
- ...
763
-
764
- @typing.overload
765
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
766
- ...
767
-
768
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
769
- """
770
- Internal decorator to support Fast bakery
771
- """
772
- ...
773
-
774
- @typing.overload
775
- 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]]]:
776
- """
777
- Creates a human-readable report, a Metaflow Card, after this step completes.
778
-
779
- Note that you may add multiple `@card` decorators in a step with different parameters.
780
-
781
-
782
- Parameters
783
- ----------
784
- type : str, default 'default'
785
- Card type.
786
- id : str, optional, default None
787
- If multiple cards are present, use this id to identify this card.
788
- options : Dict[str, Any], default {}
789
- Options passed to the card. The contents depend on the card type.
790
- timeout : int, default 45
791
- Interrupt reporting if it takes more than this many seconds.
792
- """
793
- ...
794
-
795
- @typing.overload
796
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
797
- ...
798
-
799
- @typing.overload
800
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
801
- ...
802
-
803
- 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):
804
- """
805
- Creates a human-readable report, a Metaflow Card, after this step completes.
806
-
807
- Note that you may add multiple `@card` decorators in a step with different parameters.
808
-
809
-
810
- Parameters
811
- ----------
812
- type : str, default 'default'
813
- Card type.
814
- id : str, optional, default None
815
- If multiple cards are present, use this id to identify this card.
816
- options : Dict[str, Any], default {}
817
- Options passed to the card. The contents depend on the card type.
818
- timeout : int, default 45
819
- Interrupt reporting if it takes more than this many seconds.
820
633
  """
821
634
  ...
822
635
 
@@ -909,36 +722,155 @@ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: ty
909
722
  """
910
723
  ...
911
724
 
725
+ 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]]]:
726
+ """
727
+ Specifies that this step should execute on DGX cloud.
728
+
729
+
730
+ Parameters
731
+ ----------
732
+ gpu : int
733
+ Number of GPUs to use.
734
+ gpu_type : str
735
+ Type of Nvidia GPU to use.
736
+ """
737
+ ...
738
+
912
739
  @typing.overload
913
- 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]]]:
740
+ 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]]]:
914
741
  """
915
- Specifies environment variables to be set prior to the execution of a step.
742
+ Specifies the Conda environment for the step.
743
+
744
+ Information in this decorator will augment any
745
+ attributes set in the `@conda_base` flow-level decorator. Hence,
746
+ you can use `@conda_base` to set packages required by all
747
+ steps and use `@conda` to specify step-specific overrides.
916
748
 
917
749
 
918
750
  Parameters
919
751
  ----------
920
- vars : Dict[str, str], default {}
921
- Dictionary of environment variables to set.
752
+ packages : Dict[str, str], default {}
753
+ Packages to use for this step. The key is the name of the package
754
+ and the value is the version to use.
755
+ libraries : Dict[str, str], default {}
756
+ Supported for backward compatibility. When used with packages, packages will take precedence.
757
+ python : str, optional, default None
758
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
759
+ that the version used will correspond to the version of the Python interpreter used to start the run.
760
+ disabled : bool, default False
761
+ If set to True, disables @conda.
922
762
  """
923
763
  ...
924
764
 
925
765
  @typing.overload
926
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
766
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
927
767
  ...
928
768
 
929
769
  @typing.overload
930
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
770
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
931
771
  ...
932
772
 
933
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
773
+ def conda(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
934
774
  """
935
- Specifies environment variables to be set prior to the execution of a step.
775
+ Specifies the Conda environment for the step.
776
+
777
+ Information in this decorator will augment any
778
+ attributes set in the `@conda_base` flow-level decorator. Hence,
779
+ you can use `@conda_base` to set packages required by all
780
+ steps and use `@conda` to specify step-specific overrides.
936
781
 
937
782
 
938
783
  Parameters
939
784
  ----------
940
- vars : Dict[str, str], default {}
941
- Dictionary of environment variables to set.
785
+ packages : Dict[str, str], default {}
786
+ Packages to use for this step. The key is the name of the package
787
+ and the value is the version to use.
788
+ libraries : Dict[str, str], default {}
789
+ Supported for backward compatibility. When used with packages, packages will take precedence.
790
+ python : str, optional, default None
791
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
792
+ that the version used will correspond to the version of the Python interpreter used to start the run.
793
+ disabled : bool, default False
794
+ If set to True, disables @conda.
795
+ """
796
+ ...
797
+
798
+ @typing.overload
799
+ 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]]]:
800
+ """
801
+ Specifies the resources needed when executing this step.
802
+
803
+ Use `@resources` to specify the resource requirements
804
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
805
+
806
+ You can choose the compute layer on the command line by executing e.g.
807
+ ```
808
+ python myflow.py run --with batch
809
+ ```
810
+ or
811
+ ```
812
+ python myflow.py run --with kubernetes
813
+ ```
814
+ which executes the flow on the desired system using the
815
+ requirements specified in `@resources`.
816
+
817
+
818
+ Parameters
819
+ ----------
820
+ cpu : int, default 1
821
+ Number of CPUs required for this step.
822
+ gpu : int, optional, default None
823
+ Number of GPUs required for this step.
824
+ disk : int, optional, default None
825
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
826
+ memory : int, default 4096
827
+ Memory size (in MB) required for this step.
828
+ shared_memory : int, optional, default None
829
+ The value for the size (in MiB) of the /dev/shm volume for this step.
830
+ This parameter maps to the `--shm-size` option in Docker.
831
+ """
832
+ ...
833
+
834
+ @typing.overload
835
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
836
+ ...
837
+
838
+ @typing.overload
839
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
840
+ ...
841
+
842
+ def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None):
843
+ """
844
+ Specifies the resources needed when executing this step.
845
+
846
+ Use `@resources` to specify the resource requirements
847
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
848
+
849
+ You can choose the compute layer on the command line by executing e.g.
850
+ ```
851
+ python myflow.py run --with batch
852
+ ```
853
+ or
854
+ ```
855
+ python myflow.py run --with kubernetes
856
+ ```
857
+ which executes the flow on the desired system using the
858
+ requirements specified in `@resources`.
859
+
860
+
861
+ Parameters
862
+ ----------
863
+ cpu : int, default 1
864
+ Number of CPUs required for this step.
865
+ gpu : int, optional, default None
866
+ Number of GPUs required for this step.
867
+ disk : int, optional, default None
868
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
869
+ memory : int, default 4096
870
+ Memory size (in MB) required for this step.
871
+ shared_memory : int, optional, default None
872
+ The value for the size (in MiB) of the /dev/shm volume for this step.
873
+ This parameter maps to the `--shm-size` option in Docker.
942
874
  """
943
875
  ...
944
876
 
@@ -1001,18 +933,163 @@ def timeout(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None],
1001
933
  """
1002
934
  ...
1003
935
 
1004
- 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]]]:
936
+ @typing.overload
937
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1005
938
  """
1006
- Specifies that this step is used to deploy an instance of the app.
1007
- Requires that self.app_name, self.app_port, self.entrypoint and self.deployDir is set.
939
+ Internal decorator to support Fast bakery
940
+ """
941
+ ...
942
+
943
+ @typing.overload
944
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
945
+ ...
946
+
947
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
948
+ """
949
+ Internal decorator to support Fast bakery
950
+ """
951
+ ...
952
+
953
+ def nvidia(*, gpu: int, gpu_type: str, queue_timeout: int) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
954
+ """
955
+ Specifies that this step should execute on DGX cloud.
1008
956
 
1009
957
 
1010
958
  Parameters
1011
959
  ----------
1012
- app_port : int
960
+ gpu : int
1013
961
  Number of GPUs to use.
1014
- app_name : str
1015
- Name of the app to deploy.
962
+ gpu_type : str
963
+ Type of Nvidia GPU to use.
964
+ queue_timeout : int
965
+ Time to keep the job in NVCF's queue.
966
+ """
967
+ ...
968
+
969
+ @typing.overload
970
+ 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]]]:
971
+ """
972
+ Creates a human-readable report, a Metaflow Card, after this step completes.
973
+
974
+ Note that you may add multiple `@card` decorators in a step with different parameters.
975
+
976
+
977
+ Parameters
978
+ ----------
979
+ type : str, default 'default'
980
+ Card type.
981
+ id : str, optional, default None
982
+ If multiple cards are present, use this id to identify this card.
983
+ options : Dict[str, Any], default {}
984
+ Options passed to the card. The contents depend on the card type.
985
+ timeout : int, default 45
986
+ Interrupt reporting if it takes more than this many seconds.
987
+ """
988
+ ...
989
+
990
+ @typing.overload
991
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
992
+ ...
993
+
994
+ @typing.overload
995
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
996
+ ...
997
+
998
+ 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):
999
+ """
1000
+ Creates a human-readable report, a Metaflow Card, after this step completes.
1001
+
1002
+ Note that you may add multiple `@card` decorators in a step with different parameters.
1003
+
1004
+
1005
+ Parameters
1006
+ ----------
1007
+ type : str, default 'default'
1008
+ Card type.
1009
+ id : str, optional, default None
1010
+ If multiple cards are present, use this id to identify this card.
1011
+ options : Dict[str, Any], default {}
1012
+ Options passed to the card. The contents depend on the card type.
1013
+ timeout : int, default 45
1014
+ Interrupt reporting if it takes more than this many seconds.
1015
+ """
1016
+ ...
1017
+
1018
+ def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1019
+ """
1020
+ Specifies what flows belong to the same project.
1021
+
1022
+ A project-specific namespace is created for all flows that
1023
+ use the same `@project(name)`.
1024
+
1025
+
1026
+ Parameters
1027
+ ----------
1028
+ name : str
1029
+ Project name. Make sure that the name is unique amongst all
1030
+ projects that use the same production scheduler. The name may
1031
+ contain only lowercase alphanumeric characters and underscores.
1032
+
1033
+ branch : Optional[str], default None
1034
+ The branch to use. If not specified, the branch is set to
1035
+ `user.<username>` unless `production` is set to `True`. This can
1036
+ also be set on the command line using `--branch` as a top-level option.
1037
+ It is an error to specify `branch` in the decorator and on the command line.
1038
+
1039
+ production : bool, default False
1040
+ Whether or not the branch is the production branch. This can also be set on the
1041
+ command line using `--production` as a top-level option. It is an error to specify
1042
+ `production` in the decorator and on the command line.
1043
+ The project branch name will be:
1044
+ - if `branch` is specified:
1045
+ - if `production` is True: `prod.<branch>`
1046
+ - if `production` is False: `test.<branch>`
1047
+ - if `branch` is not specified:
1048
+ - if `production` is True: `prod`
1049
+ - if `production` is False: `user.<username>`
1050
+ """
1051
+ ...
1052
+
1053
+ def airflow_external_task_sensor(*, timeout: int, poke_interval: int, mode: str, exponential_backoff: bool, pool: str, soft_fail: bool, name: str, description: str, external_dag_id: str, external_task_ids: typing.List[str], allowed_states: typing.List[str], failed_states: typing.List[str], execution_delta: "datetime.timedelta", check_existence: bool) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1054
+ """
1055
+ The `@airflow_external_task_sensor` decorator attaches a Airflow [ExternalTaskSensor](https://airflow.apache.org/docs/apache-airflow/stable/_api/airflow/sensors/external_task/index.html#airflow.sensors.external_task.ExternalTaskSensor) before the start step of the flow.
1056
+ This decorator only works when a flow is scheduled on Airflow and is compiled using `airflow create`. More than one `@airflow_external_task_sensor` can be added as a flow decorators. Adding more than one decorator will ensure that `start` step starts only after all sensors finish.
1057
+
1058
+
1059
+ Parameters
1060
+ ----------
1061
+ timeout : int
1062
+ Time, in seconds before the task times out and fails. (Default: 3600)
1063
+ poke_interval : int
1064
+ Time in seconds that the job should wait in between each try. (Default: 60)
1065
+ mode : str
1066
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1067
+ exponential_backoff : bool
1068
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1069
+ pool : str
1070
+ the slot pool this task should run in,
1071
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1072
+ soft_fail : bool
1073
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1074
+ name : str
1075
+ Name of the sensor on Airflow
1076
+ description : str
1077
+ Description of sensor in the Airflow UI
1078
+ external_dag_id : str
1079
+ The dag_id that contains the task you want to wait for.
1080
+ external_task_ids : List[str]
1081
+ The list of task_ids that you want to wait for.
1082
+ If None (default value) the sensor waits for the DAG. (Default: None)
1083
+ allowed_states : List[str]
1084
+ Iterable of allowed states, (Default: ['success'])
1085
+ failed_states : List[str]
1086
+ Iterable of failed or dis-allowed states. (Default: None)
1087
+ execution_delta : datetime.timedelta
1088
+ time difference with the previous execution to look at,
1089
+ the default is the same logical date as the current task or DAG. (Default: None)
1090
+ check_existence: bool
1091
+ Set to True to check if the external task exists or check if
1092
+ the DAG to wait for exists. (Default: True)
1016
1093
  """
1017
1094
  ...
1018
1095
 
@@ -1067,6 +1144,99 @@ def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packa
1067
1144
  """
1068
1145
  ...
1069
1146
 
1147
+ @typing.overload
1148
+ 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]]:
1149
+ """
1150
+ Specifies the event(s) that this flow depends on.
1151
+
1152
+ ```
1153
+ @trigger(event='foo')
1154
+ ```
1155
+ or
1156
+ ```
1157
+ @trigger(events=['foo', 'bar'])
1158
+ ```
1159
+
1160
+ Additionally, you can specify the parameter mappings
1161
+ to map event payload to Metaflow parameters for the flow.
1162
+ ```
1163
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1164
+ ```
1165
+ or
1166
+ ```
1167
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1168
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1169
+ ```
1170
+
1171
+ 'parameters' can also be a list of strings and tuples like so:
1172
+ ```
1173
+ @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1174
+ ```
1175
+ This is equivalent to:
1176
+ ```
1177
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1178
+ ```
1179
+
1180
+
1181
+ Parameters
1182
+ ----------
1183
+ event : Union[str, Dict[str, Any]], optional, default None
1184
+ Event dependency for this flow.
1185
+ events : List[Union[str, Dict[str, Any]]], default []
1186
+ Events dependency for this flow.
1187
+ options : Dict[str, Any], default {}
1188
+ Backend-specific configuration for tuning eventing behavior.
1189
+ """
1190
+ ...
1191
+
1192
+ @typing.overload
1193
+ def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1194
+ ...
1195
+
1196
+ 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] = {}):
1197
+ """
1198
+ Specifies the event(s) that this flow depends on.
1199
+
1200
+ ```
1201
+ @trigger(event='foo')
1202
+ ```
1203
+ or
1204
+ ```
1205
+ @trigger(events=['foo', 'bar'])
1206
+ ```
1207
+
1208
+ Additionally, you can specify the parameter mappings
1209
+ to map event payload to Metaflow parameters for the flow.
1210
+ ```
1211
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1212
+ ```
1213
+ or
1214
+ ```
1215
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1216
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1217
+ ```
1218
+
1219
+ 'parameters' can also be a list of strings and tuples like so:
1220
+ ```
1221
+ @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1222
+ ```
1223
+ This is equivalent to:
1224
+ ```
1225
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1226
+ ```
1227
+
1228
+
1229
+ Parameters
1230
+ ----------
1231
+ event : Union[str, Dict[str, Any]], optional, default None
1232
+ Event dependency for this flow.
1233
+ events : List[Union[str, Dict[str, Any]]], default []
1234
+ Events dependency for this flow.
1235
+ options : Dict[str, Any], default {}
1236
+ Backend-specific configuration for tuning eventing behavior.
1237
+ """
1238
+ ...
1239
+
1070
1240
  @typing.overload
1071
1241
  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]]:
1072
1242
  """
@@ -1157,179 +1327,78 @@ def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None)
1157
1327
 
1158
1328
  ```
1159
1329
 
1160
- - Using credentials to access the s3-compatible datastore.
1161
-
1162
- ```python
1163
- @with_artifact_store(
1164
- type="s3",
1165
- config=lambda: {
1166
- "root": "s3://my-bucket-foo/path/to/root",
1167
- "client_params": {
1168
- "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1169
- "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1170
- },
1171
- },
1172
- )
1173
- class MyFlow(FlowSpec):
1174
-
1175
- @checkpoint
1176
- @step
1177
- def start(self):
1178
- with open("my_file.txt", "w") as f:
1179
- f.write("Hello, World!")
1180
- self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1181
- self.next(self.end)
1182
-
1183
- ```
1184
-
1185
- - Accessing objects stored in external datastores after task execution.
1186
-
1187
- ```python
1188
- run = Run("CheckpointsTestsFlow/8992")
1189
- with artifact_store_from(run=run, config={
1190
- "client_params": {
1191
- "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1192
- "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1193
- },
1194
- }):
1195
- with Checkpoint() as cp:
1196
- latest = cp.list(
1197
- task=run["start"].task
1198
- )[0]
1199
- print(latest)
1200
- cp.load(
1201
- latest,
1202
- "test-checkpoints"
1203
- )
1204
-
1205
- task = Task("TorchTuneFlow/8484/train/53673")
1206
- with artifact_store_from(run=run, config={
1207
- "client_params": {
1208
- "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1209
- "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1210
- },
1211
- }):
1212
- load_model(
1213
- task.data.model_ref,
1214
- "test-models"
1215
- )
1216
- ```
1217
- Parameters:
1218
- ----------
1219
-
1220
- type: str
1221
- The type of the datastore. Can be one of 's3', 'gcs', 'azure' or any other supported metaflow Datastore.
1222
-
1223
- config: dict or Callable
1224
- Dictionary of configuration options for the datastore. The following keys are required:
1225
- - root: The root path in the datastore where the data will be saved. (needs to be in the format expected by the datastore)
1226
- - example: 's3://bucket-name/path/to/root'
1227
- - example: 'gs://bucket-name/path/to/root'
1228
- - example: 'https://myblockacc.blob.core.windows.net/metaflow/'
1229
- - role_arn (optional): AWS IAM role to access s3 bucket (only when `type` is 's3')
1230
- - session_vars (optional): AWS session variables to access s3 bucket (only when `type` is 's3')
1231
- - client_params (optional): AWS client parameters to access s3 bucket (only when `type` is 's3')
1232
- """
1233
- ...
1234
-
1235
- @typing.overload
1236
- 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]]:
1237
- """
1238
- Specifies the flow(s) that this flow depends on.
1239
-
1240
- ```
1241
- @trigger_on_finish(flow='FooFlow')
1242
- ```
1243
- or
1244
- ```
1245
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1246
- ```
1247
- This decorator respects the @project decorator and triggers the flow
1248
- when upstream runs within the same namespace complete successfully
1249
-
1250
- Additionally, you can specify project aware upstream flow dependencies
1251
- by specifying the fully qualified project_flow_name.
1252
- ```
1253
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1254
- ```
1255
- or
1256
- ```
1257
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1258
- ```
1259
-
1260
- You can also specify just the project or project branch (other values will be
1261
- inferred from the current project or project branch):
1262
- ```
1263
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1264
- ```
1265
-
1266
- Note that `branch` is typically one of:
1267
- - `prod`
1268
- - `user.bob`
1269
- - `test.my_experiment`
1270
- - `prod.staging`
1271
-
1272
-
1273
- Parameters
1274
- ----------
1275
- flow : Union[str, Dict[str, str]], optional, default None
1276
- Upstream flow dependency for this flow.
1277
- flows : List[Union[str, Dict[str, str]]], default []
1278
- Upstream flow dependencies for this flow.
1279
- options : Dict[str, Any], default {}
1280
- Backend-specific configuration for tuning eventing behavior.
1281
- """
1282
- ...
1283
-
1284
- @typing.overload
1285
- def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1286
- ...
1287
-
1288
- 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] = {}):
1289
- """
1290
- Specifies the flow(s) that this flow depends on.
1291
-
1292
- ```
1293
- @trigger_on_finish(flow='FooFlow')
1294
- ```
1295
- or
1296
- ```
1297
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1298
- ```
1299
- This decorator respects the @project decorator and triggers the flow
1300
- when upstream runs within the same namespace complete successfully
1330
+ - Using credentials to access the s3-compatible datastore.
1301
1331
 
1302
- Additionally, you can specify project aware upstream flow dependencies
1303
- by specifying the fully qualified project_flow_name.
1304
- ```
1305
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1306
- ```
1307
- or
1308
- ```
1309
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1310
- ```
1332
+ ```python
1333
+ @with_artifact_store(
1334
+ type="s3",
1335
+ config=lambda: {
1336
+ "root": "s3://my-bucket-foo/path/to/root",
1337
+ "client_params": {
1338
+ "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1339
+ "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1340
+ },
1341
+ },
1342
+ )
1343
+ class MyFlow(FlowSpec):
1311
1344
 
1312
- You can also specify just the project or project branch (other values will be
1313
- inferred from the current project or project branch):
1314
- ```
1315
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1316
- ```
1345
+ @checkpoint
1346
+ @step
1347
+ def start(self):
1348
+ with open("my_file.txt", "w") as f:
1349
+ f.write("Hello, World!")
1350
+ self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1351
+ self.next(self.end)
1317
1352
 
1318
- Note that `branch` is typically one of:
1319
- - `prod`
1320
- - `user.bob`
1321
- - `test.my_experiment`
1322
- - `prod.staging`
1353
+ ```
1323
1354
 
1355
+ - Accessing objects stored in external datastores after task execution.
1324
1356
 
1325
- Parameters
1357
+ ```python
1358
+ run = Run("CheckpointsTestsFlow/8992")
1359
+ with artifact_store_from(run=run, config={
1360
+ "client_params": {
1361
+ "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1362
+ "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1363
+ },
1364
+ }):
1365
+ with Checkpoint() as cp:
1366
+ latest = cp.list(
1367
+ task=run["start"].task
1368
+ )[0]
1369
+ print(latest)
1370
+ cp.load(
1371
+ latest,
1372
+ "test-checkpoints"
1373
+ )
1374
+
1375
+ task = Task("TorchTuneFlow/8484/train/53673")
1376
+ with artifact_store_from(run=run, config={
1377
+ "client_params": {
1378
+ "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1379
+ "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1380
+ },
1381
+ }):
1382
+ load_model(
1383
+ task.data.model_ref,
1384
+ "test-models"
1385
+ )
1386
+ ```
1387
+ Parameters:
1326
1388
  ----------
1327
- flow : Union[str, Dict[str, str]], optional, default None
1328
- Upstream flow dependency for this flow.
1329
- flows : List[Union[str, Dict[str, str]]], default []
1330
- Upstream flow dependencies for this flow.
1331
- options : Dict[str, Any], default {}
1332
- Backend-specific configuration for tuning eventing behavior.
1389
+
1390
+ type: str
1391
+ The type of the datastore. Can be one of 's3', 'gcs', 'azure' or any other supported metaflow Datastore.
1392
+
1393
+ config: dict or Callable
1394
+ Dictionary of configuration options for the datastore. The following keys are required:
1395
+ - root: The root path in the datastore where the data will be saved. (needs to be in the format expected by the datastore)
1396
+ - example: 's3://bucket-name/path/to/root'
1397
+ - example: 'gs://bucket-name/path/to/root'
1398
+ - example: 'https://myblockacc.blob.core.windows.net/metaflow/'
1399
+ - role_arn (optional): AWS IAM role to access s3 bucket (only when `type` is 's3')
1400
+ - session_vars (optional): AWS session variables to access s3 bucket (only when `type` is 's3')
1401
+ - client_params (optional): AWS client parameters to access s3 bucket (only when `type` is 's3')
1333
1402
  """
1334
1403
  ...
1335
1404
 
@@ -1376,84 +1445,6 @@ def airflow_s3_key_sensor(*, timeout: int, poke_interval: int, mode: str, expone
1376
1445
  """
1377
1446
  ...
1378
1447
 
1379
- def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1380
- """
1381
- Specifies what flows belong to the same project.
1382
-
1383
- A project-specific namespace is created for all flows that
1384
- use the same `@project(name)`.
1385
-
1386
-
1387
- Parameters
1388
- ----------
1389
- name : str
1390
- Project name. Make sure that the name is unique amongst all
1391
- projects that use the same production scheduler. The name may
1392
- contain only lowercase alphanumeric characters and underscores.
1393
-
1394
- branch : Optional[str], default None
1395
- The branch to use. If not specified, the branch is set to
1396
- `user.<username>` unless `production` is set to `True`. This can
1397
- also be set on the command line using `--branch` as a top-level option.
1398
- It is an error to specify `branch` in the decorator and on the command line.
1399
-
1400
- production : bool, default False
1401
- Whether or not the branch is the production branch. This can also be set on the
1402
- command line using `--production` as a top-level option. It is an error to specify
1403
- `production` in the decorator and on the command line.
1404
- The project branch name will be:
1405
- - if `branch` is specified:
1406
- - if `production` is True: `prod.<branch>`
1407
- - if `production` is False: `test.<branch>`
1408
- - if `branch` is not specified:
1409
- - if `production` is True: `prod`
1410
- - if `production` is False: `user.<username>`
1411
- """
1412
- ...
1413
-
1414
- def airflow_external_task_sensor(*, timeout: int, poke_interval: int, mode: str, exponential_backoff: bool, pool: str, soft_fail: bool, name: str, description: str, external_dag_id: str, external_task_ids: typing.List[str], allowed_states: typing.List[str], failed_states: typing.List[str], execution_delta: "datetime.timedelta", check_existence: bool) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1415
- """
1416
- The `@airflow_external_task_sensor` decorator attaches a Airflow [ExternalTaskSensor](https://airflow.apache.org/docs/apache-airflow/stable/_api/airflow/sensors/external_task/index.html#airflow.sensors.external_task.ExternalTaskSensor) before the start step of the flow.
1417
- This decorator only works when a flow is scheduled on Airflow and is compiled using `airflow create`. More than one `@airflow_external_task_sensor` can be added as a flow decorators. Adding more than one decorator will ensure that `start` step starts only after all sensors finish.
1418
-
1419
-
1420
- Parameters
1421
- ----------
1422
- timeout : int
1423
- Time, in seconds before the task times out and fails. (Default: 3600)
1424
- poke_interval : int
1425
- Time in seconds that the job should wait in between each try. (Default: 60)
1426
- mode : str
1427
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1428
- exponential_backoff : bool
1429
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1430
- pool : str
1431
- the slot pool this task should run in,
1432
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1433
- soft_fail : bool
1434
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1435
- name : str
1436
- Name of the sensor on Airflow
1437
- description : str
1438
- Description of sensor in the Airflow UI
1439
- external_dag_id : str
1440
- The dag_id that contains the task you want to wait for.
1441
- external_task_ids : List[str]
1442
- The list of task_ids that you want to wait for.
1443
- If None (default value) the sensor waits for the DAG. (Default: None)
1444
- allowed_states : List[str]
1445
- Iterable of allowed states, (Default: ['success'])
1446
- failed_states : List[str]
1447
- Iterable of failed or dis-allowed states. (Default: None)
1448
- execution_delta : datetime.timedelta
1449
- time difference with the previous execution to look at,
1450
- the default is the same logical date as the current task or DAG. (Default: None)
1451
- check_existence: bool
1452
- Set to True to check if the external task exists or check if
1453
- the DAG to wait for exists. (Default: True)
1454
- """
1455
- ...
1456
-
1457
1448
  @typing.overload
1458
1449
  def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1459
1450
  """
@@ -1496,93 +1487,101 @@ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packag
1496
1487
  ...
1497
1488
 
1498
1489
  @typing.overload
1499
- 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]]:
1490
+ 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]]:
1500
1491
  """
1501
- Specifies the event(s) that this flow depends on.
1492
+ Specifies the flow(s) that this flow depends on.
1502
1493
 
1503
1494
  ```
1504
- @trigger(event='foo')
1495
+ @trigger_on_finish(flow='FooFlow')
1505
1496
  ```
1506
1497
  or
1507
1498
  ```
1508
- @trigger(events=['foo', 'bar'])
1499
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1509
1500
  ```
1501
+ This decorator respects the @project decorator and triggers the flow
1502
+ when upstream runs within the same namespace complete successfully
1510
1503
 
1511
- Additionally, you can specify the parameter mappings
1512
- to map event payload to Metaflow parameters for the flow.
1504
+ Additionally, you can specify project aware upstream flow dependencies
1505
+ by specifying the fully qualified project_flow_name.
1513
1506
  ```
1514
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1507
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1515
1508
  ```
1516
1509
  or
1517
1510
  ```
1518
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1519
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1511
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1520
1512
  ```
1521
1513
 
1522
- 'parameters' can also be a list of strings and tuples like so:
1523
- ```
1524
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1525
- ```
1526
- This is equivalent to:
1514
+ You can also specify just the project or project branch (other values will be
1515
+ inferred from the current project or project branch):
1527
1516
  ```
1528
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1517
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1529
1518
  ```
1530
1519
 
1520
+ Note that `branch` is typically one of:
1521
+ - `prod`
1522
+ - `user.bob`
1523
+ - `test.my_experiment`
1524
+ - `prod.staging`
1525
+
1531
1526
 
1532
1527
  Parameters
1533
1528
  ----------
1534
- event : Union[str, Dict[str, Any]], optional, default None
1535
- Event dependency for this flow.
1536
- events : List[Union[str, Dict[str, Any]]], default []
1537
- Events dependency for this flow.
1529
+ flow : Union[str, Dict[str, str]], optional, default None
1530
+ Upstream flow dependency for this flow.
1531
+ flows : List[Union[str, Dict[str, str]]], default []
1532
+ Upstream flow dependencies for this flow.
1538
1533
  options : Dict[str, Any], default {}
1539
1534
  Backend-specific configuration for tuning eventing behavior.
1540
1535
  """
1541
1536
  ...
1542
1537
 
1543
1538
  @typing.overload
1544
- def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1539
+ def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1545
1540
  ...
1546
1541
 
1547
- 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] = {}):
1542
+ 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] = {}):
1548
1543
  """
1549
- Specifies the event(s) that this flow depends on.
1544
+ Specifies the flow(s) that this flow depends on.
1550
1545
 
1551
1546
  ```
1552
- @trigger(event='foo')
1547
+ @trigger_on_finish(flow='FooFlow')
1553
1548
  ```
1554
1549
  or
1555
1550
  ```
1556
- @trigger(events=['foo', 'bar'])
1551
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1557
1552
  ```
1553
+ This decorator respects the @project decorator and triggers the flow
1554
+ when upstream runs within the same namespace complete successfully
1558
1555
 
1559
- Additionally, you can specify the parameter mappings
1560
- to map event payload to Metaflow parameters for the flow.
1556
+ Additionally, you can specify project aware upstream flow dependencies
1557
+ by specifying the fully qualified project_flow_name.
1561
1558
  ```
1562
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1559
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1563
1560
  ```
1564
1561
  or
1565
1562
  ```
1566
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1567
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1563
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1568
1564
  ```
1569
1565
 
1570
- 'parameters' can also be a list of strings and tuples like so:
1571
- ```
1572
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1573
- ```
1574
- This is equivalent to:
1566
+ You can also specify just the project or project branch (other values will be
1567
+ inferred from the current project or project branch):
1575
1568
  ```
1576
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1569
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1577
1570
  ```
1578
1571
 
1572
+ Note that `branch` is typically one of:
1573
+ - `prod`
1574
+ - `user.bob`
1575
+ - `test.my_experiment`
1576
+ - `prod.staging`
1577
+
1579
1578
 
1580
1579
  Parameters
1581
1580
  ----------
1582
- event : Union[str, Dict[str, Any]], optional, default None
1583
- Event dependency for this flow.
1584
- events : List[Union[str, Dict[str, Any]]], default []
1585
- Events dependency for this flow.
1581
+ flow : Union[str, Dict[str, str]], optional, default None
1582
+ Upstream flow dependency for this flow.
1583
+ flows : List[Union[str, Dict[str, str]]], default []
1584
+ Upstream flow dependencies for this flow.
1586
1585
  options : Dict[str, Any], default {}
1587
1586
  Backend-specific configuration for tuning eventing behavior.
1588
1587
  """