ob-metaflow-stubs 6.0.3.185__py2.py3-none-any.whl → 6.0.3.187__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 (216) hide show
  1. metaflow-stubs/__init__.pyi +1086 -857
  2. metaflow-stubs/cards.pyi +2 -2
  3. metaflow-stubs/cli.pyi +2 -2
  4. metaflow-stubs/cli_components/__init__.pyi +2 -2
  5. metaflow-stubs/cli_components/utils.pyi +2 -2
  6. metaflow-stubs/client/__init__.pyi +2 -2
  7. metaflow-stubs/client/core.pyi +5 -5
  8. metaflow-stubs/client/filecache.pyi +2 -2
  9. metaflow-stubs/events.pyi +3 -3
  10. metaflow-stubs/exception.pyi +2 -2
  11. metaflow-stubs/flowspec.pyi +6 -6
  12. metaflow-stubs/generated_for.txt +1 -1
  13. metaflow-stubs/includefile.pyi +4 -4
  14. metaflow-stubs/info_file.pyi +2 -2
  15. metaflow-stubs/metadata_provider/__init__.pyi +2 -2
  16. metaflow-stubs/metadata_provider/heartbeat.pyi +2 -2
  17. metaflow-stubs/metadata_provider/metadata.pyi +3 -3
  18. metaflow-stubs/metadata_provider/util.pyi +2 -2
  19. metaflow-stubs/metaflow_config.pyi +2 -2
  20. metaflow-stubs/metaflow_current.pyi +53 -201
  21. metaflow-stubs/metaflow_git.pyi +2 -2
  22. metaflow-stubs/mf_extensions/__init__.pyi +2 -2
  23. metaflow-stubs/mf_extensions/obcheckpoint/__init__.pyi +2 -2
  24. metaflow-stubs/mf_extensions/obcheckpoint/plugins/__init__.pyi +2 -2
  25. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/__init__.pyi +2 -2
  26. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/__init__.pyi +2 -2
  27. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/async_cards.pyi +2 -2
  28. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/deco_injection_mixin.pyi +2 -2
  29. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/extra_components.pyi +3 -3
  30. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/__init__.pyi +2 -2
  31. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/__init__.pyi +2 -2
  32. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/checkpoint_lister.pyi +3 -3
  33. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/lineage_card.pyi +2 -2
  34. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/checkpoint_storage.pyi +4 -4
  35. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/constructors.pyi +2 -2
  36. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/core.pyi +4 -4
  37. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/decorator.pyi +51 -49
  38. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/exceptions.pyi +2 -2
  39. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/final_api.pyi +11 -15
  40. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/lineage.pyi +2 -2
  41. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/__init__.pyi +2 -2
  42. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/context.pyi +3 -3
  43. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/core.pyi +2 -2
  44. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/decorator.pyi +2 -2
  45. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/exceptions.pyi +2 -2
  46. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/task_utils.pyi +4 -4
  47. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/utils.pyi +2 -2
  48. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastructures.pyi +87 -3
  49. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/exceptions.pyi +2 -2
  50. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/__init__.pyi +2 -2
  51. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/decorator.pyi +79 -76
  52. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/__init__.pyi +2 -2
  53. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/core.pyi +48 -13
  54. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/exceptions.pyi +2 -2
  55. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/model_storage.pyi +3 -3
  56. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/__init__.pyi +2 -2
  57. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/flowspec_utils.pyi +2 -2
  58. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/general.pyi +2 -2
  59. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/identity_utils.pyi +3 -3
  60. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/__init__.pyi +2 -2
  61. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/base.pyi +2 -2
  62. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/tar.pyi +3 -3
  63. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/tar_utils.pyi +3 -3
  64. metaflow-stubs/mf_extensions/outerbounds/__init__.pyi +2 -2
  65. metaflow-stubs/mf_extensions/outerbounds/plugins/__init__.pyi +2 -2
  66. metaflow-stubs/mf_extensions/outerbounds/plugins/aws/__init__.pyi +2 -2
  67. metaflow-stubs/mf_extensions/outerbounds/plugins/aws/assume_role_decorator.pyi +2 -2
  68. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/__init__.pyi +2 -2
  69. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/async_cards.pyi +2 -2
  70. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/injector.pyi +2 -2
  71. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/__init__.pyi +2 -2
  72. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/coreweave.pyi +2 -2
  73. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/nebius.pyi +2 -2
  74. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/__init__.pyi +2 -2
  75. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/pod_killer.pyi +2 -2
  76. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/__init__.pyi +2 -2
  77. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/constants.pyi +2 -2
  78. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/exceptions.pyi +2 -2
  79. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/ollama.pyi +2 -2
  80. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/status_card.pyi +2 -2
  81. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/__init__.pyi +2 -2
  82. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/snowflake.pyi +2 -2
  83. metaflow-stubs/mf_extensions/outerbounds/profilers/__init__.pyi +2 -2
  84. metaflow-stubs/mf_extensions/outerbounds/profilers/gpu.pyi +2 -2
  85. metaflow-stubs/mf_extensions/outerbounds/remote_config.pyi +2 -2
  86. metaflow-stubs/mf_extensions/outerbounds/toplevel/__init__.pyi +2 -2
  87. metaflow-stubs/mf_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.pyi +2 -2
  88. metaflow-stubs/multicore_utils.pyi +2 -2
  89. metaflow-stubs/ob_internal.pyi +2 -2
  90. metaflow-stubs/parameters.pyi +4 -4
  91. metaflow-stubs/plugins/__init__.pyi +12 -12
  92. metaflow-stubs/plugins/airflow/__init__.pyi +2 -2
  93. metaflow-stubs/plugins/airflow/airflow_utils.pyi +2 -2
  94. metaflow-stubs/plugins/airflow/exception.pyi +2 -2
  95. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +2 -2
  96. metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi +2 -2
  97. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +2 -2
  98. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +2 -2
  99. metaflow-stubs/plugins/argo/__init__.pyi +2 -2
  100. metaflow-stubs/plugins/argo/argo_client.pyi +2 -2
  101. metaflow-stubs/plugins/argo/argo_events.pyi +2 -2
  102. metaflow-stubs/plugins/argo/argo_workflows.pyi +3 -3
  103. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +3 -3
  104. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +4 -4
  105. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +4 -4
  106. metaflow-stubs/plugins/aws/__init__.pyi +2 -2
  107. metaflow-stubs/plugins/aws/aws_client.pyi +2 -2
  108. metaflow-stubs/plugins/aws/aws_utils.pyi +2 -2
  109. metaflow-stubs/plugins/aws/batch/__init__.pyi +2 -2
  110. metaflow-stubs/plugins/aws/batch/batch.pyi +2 -2
  111. metaflow-stubs/plugins/aws/batch/batch_client.pyi +2 -2
  112. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +2 -2
  113. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +2 -2
  114. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +5 -5
  115. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +2 -2
  116. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +2 -2
  117. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +2 -2
  118. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +2 -2
  119. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +2 -2
  120. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer.pyi +5 -5
  121. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer_objects.pyi +4 -4
  122. metaflow-stubs/plugins/azure/__init__.pyi +2 -2
  123. metaflow-stubs/plugins/azure/azure_credential.pyi +2 -2
  124. metaflow-stubs/plugins/azure/azure_exceptions.pyi +2 -2
  125. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +5 -5
  126. metaflow-stubs/plugins/azure/azure_utils.pyi +2 -2
  127. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +2 -2
  128. metaflow-stubs/plugins/azure/includefile_support.pyi +2 -2
  129. metaflow-stubs/plugins/cards/__init__.pyi +6 -6
  130. metaflow-stubs/plugins/cards/card_client.pyi +2 -2
  131. metaflow-stubs/plugins/cards/card_creator.pyi +2 -2
  132. metaflow-stubs/plugins/cards/card_datastore.pyi +2 -2
  133. metaflow-stubs/plugins/cards/card_decorator.pyi +2 -2
  134. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +2 -2
  135. metaflow-stubs/plugins/cards/card_modules/basic.pyi +2 -2
  136. metaflow-stubs/plugins/cards/card_modules/card.pyi +2 -2
  137. metaflow-stubs/plugins/cards/card_modules/components.pyi +3 -3
  138. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +2 -2
  139. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +2 -2
  140. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +2 -2
  141. metaflow-stubs/plugins/cards/card_resolver.pyi +2 -2
  142. metaflow-stubs/plugins/cards/component_serializer.pyi +2 -2
  143. metaflow-stubs/plugins/cards/exception.pyi +2 -2
  144. metaflow-stubs/plugins/catch_decorator.pyi +2 -2
  145. metaflow-stubs/plugins/datatools/__init__.pyi +2 -2
  146. metaflow-stubs/plugins/datatools/local.pyi +2 -2
  147. metaflow-stubs/plugins/datatools/s3/__init__.pyi +2 -2
  148. metaflow-stubs/plugins/datatools/s3/s3.pyi +3 -3
  149. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +2 -2
  150. metaflow-stubs/plugins/datatools/s3/s3util.pyi +2 -2
  151. metaflow-stubs/plugins/debug_logger.pyi +2 -2
  152. metaflow-stubs/plugins/debug_monitor.pyi +2 -2
  153. metaflow-stubs/plugins/environment_decorator.pyi +2 -2
  154. metaflow-stubs/plugins/events_decorator.pyi +2 -2
  155. metaflow-stubs/plugins/frameworks/__init__.pyi +2 -2
  156. metaflow-stubs/plugins/frameworks/pytorch.pyi +2 -2
  157. metaflow-stubs/plugins/gcp/__init__.pyi +2 -2
  158. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +5 -5
  159. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +2 -2
  160. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +2 -2
  161. metaflow-stubs/plugins/gcp/gs_utils.pyi +2 -2
  162. metaflow-stubs/plugins/gcp/includefile_support.pyi +2 -2
  163. metaflow-stubs/plugins/kubernetes/__init__.pyi +2 -2
  164. metaflow-stubs/plugins/kubernetes/kube_utils.pyi +3 -3
  165. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +2 -2
  166. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +2 -2
  167. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +2 -2
  168. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +2 -2
  169. metaflow-stubs/plugins/kubernetes/spot_monitor_sidecar.pyi +2 -2
  170. metaflow-stubs/plugins/ollama/__init__.pyi +3 -3
  171. metaflow-stubs/plugins/parallel_decorator.pyi +2 -2
  172. metaflow-stubs/plugins/perimeters.pyi +2 -2
  173. metaflow-stubs/plugins/project_decorator.pyi +2 -2
  174. metaflow-stubs/plugins/pypi/__init__.pyi +3 -3
  175. metaflow-stubs/plugins/pypi/conda_decorator.pyi +2 -2
  176. metaflow-stubs/plugins/pypi/conda_environment.pyi +5 -5
  177. metaflow-stubs/plugins/pypi/parsers.pyi +2 -2
  178. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +2 -2
  179. metaflow-stubs/plugins/pypi/pypi_environment.pyi +2 -2
  180. metaflow-stubs/plugins/pypi/utils.pyi +2 -2
  181. metaflow-stubs/plugins/resources_decorator.pyi +2 -2
  182. metaflow-stubs/plugins/retry_decorator.pyi +2 -2
  183. metaflow-stubs/plugins/secrets/__init__.pyi +3 -3
  184. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +4 -4
  185. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +2 -2
  186. metaflow-stubs/plugins/snowflake/__init__.pyi +2 -2
  187. metaflow-stubs/plugins/storage_executor.pyi +2 -2
  188. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +2 -2
  189. metaflow-stubs/plugins/timeout_decorator.pyi +2 -2
  190. metaflow-stubs/plugins/torchtune/__init__.pyi +2 -2
  191. metaflow-stubs/plugins/uv/__init__.pyi +2 -2
  192. metaflow-stubs/plugins/uv/uv_environment.pyi +3 -3
  193. metaflow-stubs/profilers/__init__.pyi +2 -2
  194. metaflow-stubs/pylint_wrapper.pyi +2 -2
  195. metaflow-stubs/runner/__init__.pyi +2 -2
  196. metaflow-stubs/runner/deployer.pyi +31 -31
  197. metaflow-stubs/runner/deployer_impl.pyi +3 -3
  198. metaflow-stubs/runner/metaflow_runner.pyi +3 -3
  199. metaflow-stubs/runner/nbdeploy.pyi +2 -2
  200. metaflow-stubs/runner/nbrun.pyi +2 -2
  201. metaflow-stubs/runner/subprocess_manager.pyi +2 -2
  202. metaflow-stubs/runner/utils.pyi +3 -3
  203. metaflow-stubs/system/__init__.pyi +2 -2
  204. metaflow-stubs/system/system_logger.pyi +3 -3
  205. metaflow-stubs/system/system_monitor.pyi +2 -2
  206. metaflow-stubs/tagging_util.pyi +2 -2
  207. metaflow-stubs/tuple_util.pyi +2 -2
  208. metaflow-stubs/user_configs/__init__.pyi +2 -2
  209. metaflow-stubs/user_configs/config_decorators.pyi +5 -5
  210. metaflow-stubs/user_configs/config_options.pyi +3 -3
  211. metaflow-stubs/user_configs/config_parameters.pyi +5 -5
  212. {ob_metaflow_stubs-6.0.3.185.dist-info → ob_metaflow_stubs-6.0.3.187.dist-info}/METADATA +1 -1
  213. ob_metaflow_stubs-6.0.3.187.dist-info/RECORD +216 -0
  214. ob_metaflow_stubs-6.0.3.185.dist-info/RECORD +0 -216
  215. {ob_metaflow_stubs-6.0.3.185.dist-info → ob_metaflow_stubs-6.0.3.187.dist-info}/WHEEL +0 -0
  216. {ob_metaflow_stubs-6.0.3.185.dist-info → ob_metaflow_stubs-6.0.3.187.dist-info}/top_level.txt +0 -0
@@ -1,7 +1,7 @@
1
1
  ######################################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
- # MF version: 2.15.18.1+obcheckpoint(0.2.1);ob(v1) #
4
- # Generated on 2025-06-26T22:38:03.208016 #
3
+ # MF version: 2.15.18.1+obcheckpoint(0.2.4);ob(v1) #
4
+ # Generated on 2025-07-01T15:21:03.639611 #
5
5
  ######################################################################################################
6
6
 
7
7
  from __future__ import annotations
@@ -35,18 +35,18 @@ from .user_configs.config_parameters import ConfigValue as ConfigValue
35
35
  from .user_configs.config_parameters import config_expr as config_expr
36
36
  from .user_configs.config_decorators import CustomFlowDecorator as CustomFlowDecorator
37
37
  from .user_configs.config_decorators import CustomStepDecorator as CustomStepDecorator
38
+ from . import events as events
38
39
  from . import cards as cards
39
- from . import metaflow_git as metaflow_git
40
40
  from . import tuple_util as tuple_util
41
- from . import events as events
41
+ from . import metaflow_git as metaflow_git
42
42
  from . import runner as runner
43
43
  from . import plugins as plugins
44
44
  from .mf_extensions.outerbounds.toplevel.global_aliases_for_metaflow_package import S3 as S3
45
45
  from . import includefile as includefile
46
46
  from .includefile import IncludeFile as IncludeFile
47
- from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
48
47
  from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
49
48
  from .plugins.pypi.parsers import conda_environment_yml_parser as conda_environment_yml_parser
49
+ from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
50
50
  from . import client as client
51
51
  from .client.core import namespace as namespace
52
52
  from .client.core import get_namespace as get_namespace
@@ -155,126 +155,66 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
155
155
  """
156
156
  ...
157
157
 
158
- 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]]]:
158
+ @typing.overload
159
+ 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]]]:
159
160
  """
160
- Specifies that this step should execute on DGX cloud.
161
-
161
+ Creates a human-readable report, a Metaflow Card, after this step completes.
162
162
 
163
- Parameters
164
- ----------
165
- gpu : int
166
- Number of GPUs to use.
167
- gpu_type : str
168
- Type of Nvidia GPU to use.
169
- """
170
- ...
171
-
172
- 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]]]:
173
- """
174
- Specifies that this step is used to deploy an instance of the app.
175
- Requires that self.app_name, self.app_port, self.entrypoint and self.deployDir is set.
163
+ Note that you may add multiple `@card` decorators in a step with different parameters.
176
164
 
177
165
 
178
166
  Parameters
179
167
  ----------
180
- app_port : int
181
- Number of GPUs to use.
182
- app_name : str
183
- Name of the app to deploy.
168
+ type : str, default 'default'
169
+ Card type.
170
+ id : str, optional, default None
171
+ If multiple cards are present, use this id to identify this card.
172
+ options : Dict[str, Any], default {}
173
+ Options passed to the card. The contents depend on the card type.
174
+ timeout : int, default 45
175
+ Interrupt reporting if it takes more than this many seconds.
184
176
  """
185
177
  ...
186
178
 
187
- 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]]]:
188
- """
189
- This decorator is used to run Ollama APIs as Metaflow task sidecars.
190
-
191
- User code call
192
- --------------
193
- @ollama(
194
- models=[...],
195
- ...
196
- )
197
-
198
- Valid backend options
199
- ---------------------
200
- - 'local': Run as a separate process on the local task machine.
201
- - (TODO) 'managed': Outerbounds hosts and selects compute provider.
202
- - (TODO) 'remote': Spin up separate instance to serve Ollama models.
203
-
204
- Valid model options
205
- -------------------
206
- Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
207
-
208
-
209
- Parameters
210
- ----------
211
- models: list[str]
212
- List of Ollama containers running models in sidecars.
213
- backend: str
214
- Determines where and how to run the Ollama process.
215
- force_pull: bool
216
- Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
217
- cache_update_policy: str
218
- Cache update policy: "auto", "force", or "never".
219
- force_cache_update: bool
220
- Simple override for "force" cache update policy.
221
- debug: bool
222
- Whether to turn on verbose debugging logs.
223
- circuit_breaker_config: dict
224
- Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
225
- timeout_config: dict
226
- Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
227
- """
179
+ @typing.overload
180
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
228
181
  ...
229
182
 
230
183
  @typing.overload
231
- 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]]]:
184
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
185
+ ...
186
+
187
+ 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):
232
188
  """
233
- Specifies the PyPI packages for the step.
189
+ Creates a human-readable report, a Metaflow Card, after this step completes.
234
190
 
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.
191
+ Note that you may add multiple `@card` decorators in a step with different parameters.
239
192
 
240
193
 
241
194
  Parameters
242
195
  ----------
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.
196
+ type : str, default 'default'
197
+ Card type.
198
+ id : str, optional, default None
199
+ If multiple cards are present, use this id to identify this card.
200
+ options : Dict[str, Any], default {}
201
+ Options passed to the card. The contents depend on the card type.
202
+ timeout : int, default 45
203
+ Interrupt reporting if it takes more than this many seconds.
249
204
  """
250
205
  ...
251
206
 
252
- @typing.overload
253
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
254
- ...
255
-
256
- @typing.overload
257
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
258
- ...
259
-
260
- 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):
207
+ 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]]]:
261
208
  """
262
- Specifies the PyPI packages for the step.
263
-
264
- Information in this decorator will augment any
265
- attributes set in the `@pyi_base` flow-level decorator. Hence,
266
- you can use `@pypi_base` to set packages required by all
267
- steps and use `@pypi` to specify step-specific overrides.
209
+ Specifies that this step should execute on DGX cloud.
268
210
 
269
211
 
270
212
  Parameters
271
213
  ----------
272
- packages : Dict[str, str], default: {}
273
- Packages to use for this step. The key is the name of the package
274
- and the value is the version to use.
275
- python : str, optional, default: None
276
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
277
- that the version used will correspond to the version of the Python interpreter used to start the run.
214
+ gpu : int
215
+ Number of GPUs to use.
216
+ gpu_type : str
217
+ Type of Nvidia GPU to use.
278
218
  """
279
219
  ...
280
220
 
@@ -329,306 +269,374 @@ def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
329
269
  """
330
270
  ...
331
271
 
332
- @typing.overload
333
- 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]]]:
272
+ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: typing.Optional[str] = None, image_pull_policy: str = 'KUBERNETES_IMAGE_PULL_POLICY', image_pull_secrets: typing.List[str] = [], service_account: str = 'METAFLOW_KUBERNETES_SERVICE_ACCOUNT', secrets: typing.Optional[typing.List[str]] = None, node_selector: typing.Union[typing.Dict[str, str], str, None] = None, namespace: str = 'METAFLOW_KUBERNETES_NAMESPACE', gpu: typing.Optional[int] = None, gpu_vendor: str = 'KUBERNETES_GPU_VENDOR', tolerations: typing.List[str] = [], labels: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_LABELS', annotations: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_ANNOTATIONS', use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = '/metaflow_temp', persistent_volume_claims: typing.Optional[typing.Dict[str, str]] = None, shared_memory: typing.Optional[int] = None, port: typing.Optional[int] = None, compute_pool: typing.Optional[str] = None, hostname_resolution_timeout: int = 600, qos: str = 'Burstable', security_context: typing.Optional[typing.Dict[str, typing.Any]] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
334
273
  """
335
- Enables loading / saving of models within a step.
336
-
274
+ Specifies that this step should execute on Kubernetes.
337
275
 
338
276
 
339
277
  Parameters
340
278
  ----------
341
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
342
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
343
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
344
- - `current.checkpoint`
345
- - `current.model`
346
- - `current.huggingface_hub`
347
-
348
- 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
349
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
350
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
279
+ cpu : int, default 1
280
+ Number of CPUs required for this step. If `@resources` is
281
+ also present, the maximum value from all decorators is used.
282
+ memory : int, default 4096
283
+ Memory size (in MB) required for this step. If
284
+ `@resources` is also present, the maximum value from all decorators is
285
+ used.
286
+ disk : int, default 10240
287
+ Disk size (in MB) required for this step. If
288
+ `@resources` is also present, the maximum value from all decorators is
289
+ used.
290
+ image : str, optional, default None
291
+ Docker image to use when launching on Kubernetes. If not specified, and
292
+ METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
293
+ not, a default Docker image mapping to the current version of Python is used.
294
+ image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
295
+ If given, the imagePullPolicy to be applied to the Docker image of the step.
296
+ image_pull_secrets: List[str], default []
297
+ The default is extracted from METAFLOW_KUBERNETES_IMAGE_PULL_SECRETS.
298
+ Kubernetes image pull secrets to use when pulling container images
299
+ in Kubernetes.
300
+ service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
301
+ Kubernetes service account to use when launching pod in Kubernetes.
302
+ secrets : List[str], optional, default None
303
+ Kubernetes secrets to use when launching pod in Kubernetes. These
304
+ secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
305
+ in Metaflow configuration.
306
+ node_selector: Union[Dict[str,str], str], optional, default None
307
+ Kubernetes node selector(s) to apply to the pod running the task.
308
+ Can be passed in as a comma separated string of values e.g.
309
+ 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
310
+ {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
311
+ namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
312
+ Kubernetes namespace to use when launching pod in Kubernetes.
313
+ gpu : int, optional, default None
314
+ Number of GPUs required for this step. A value of zero implies that
315
+ the scheduled node should not have GPUs.
316
+ gpu_vendor : str, default KUBERNETES_GPU_VENDOR
317
+ The vendor of the GPUs to be used for this step.
318
+ tolerations : List[str], default []
319
+ The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
320
+ Kubernetes tolerations to use when launching pod in Kubernetes.
321
+ labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
322
+ Kubernetes labels to use when launching pod in Kubernetes.
323
+ annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
324
+ Kubernetes annotations to use when launching pod in Kubernetes.
325
+ use_tmpfs : bool, default False
326
+ This enables an explicit tmpfs mount for this step.
327
+ tmpfs_tempdir : bool, default True
328
+ sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
329
+ tmpfs_size : int, optional, default: None
330
+ The value for the size (in MiB) of the tmpfs mount for this step.
331
+ This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
332
+ memory allocated for this step.
333
+ tmpfs_path : str, optional, default /metaflow_temp
334
+ Path to tmpfs mount for this step.
335
+ persistent_volume_claims : Dict[str, str], optional, default None
336
+ A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
337
+ volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
338
+ shared_memory: int, optional
339
+ Shared memory size (in MiB) required for this step
340
+ port: int, optional
341
+ Port number to specify in the Kubernetes job object
342
+ compute_pool : str, optional, default None
343
+ Compute pool to be used for for this step.
344
+ If not specified, any accessible compute pool within the perimeter is used.
345
+ hostname_resolution_timeout: int, default 10 * 60
346
+ Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
347
+ Only applicable when @parallel is used.
348
+ qos: str, default: Burstable
349
+ Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
351
350
 
352
- temp_dir_root : str, default: None
353
- The root directory under which `current.model.loaded` will store loaded models
351
+ security_context: Dict[str, Any], optional, default None
352
+ Container security context. Applies to the task container. Allows the following keys:
353
+ - privileged: bool, optional, default None
354
+ - allow_privilege_escalation: bool, optional, default None
355
+ - run_as_user: int, optional, default None
356
+ - run_as_group: int, optional, default None
357
+ - run_as_non_root: bool, optional, default None
354
358
  """
355
359
  ...
356
360
 
357
- @typing.overload
358
- def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
359
- ...
360
-
361
- @typing.overload
362
- def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
363
- ...
364
-
365
- 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):
361
+ 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]]]:
366
362
  """
367
- Enables loading / saving of models within a step.
368
-
363
+ Specifies that this step should execute on DGX cloud.
369
364
 
370
365
 
371
366
  Parameters
372
367
  ----------
373
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
374
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
375
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
376
- - `current.checkpoint`
377
- - `current.model`
378
- - `current.huggingface_hub`
379
-
380
- 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
381
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
382
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
383
-
384
- temp_dir_root : str, default: None
385
- The root directory under which `current.model.loaded` will store loaded models
368
+ gpu : int
369
+ Number of GPUs to use.
370
+ gpu_type : str
371
+ Type of Nvidia GPU to use.
372
+ queue_timeout : int
373
+ Time to keep the job in NVCF's queue.
386
374
  """
387
375
  ...
388
376
 
389
377
  @typing.overload
390
- 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]]]:
378
+ 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]]]:
391
379
  """
392
- Specifies environment variables to be set prior to the execution of a step.
380
+ Specifies the resources needed when executing this step.
381
+
382
+ Use `@resources` to specify the resource requirements
383
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
384
+
385
+ You can choose the compute layer on the command line by executing e.g.
386
+ ```
387
+ python myflow.py run --with batch
388
+ ```
389
+ or
390
+ ```
391
+ python myflow.py run --with kubernetes
392
+ ```
393
+ which executes the flow on the desired system using the
394
+ requirements specified in `@resources`.
393
395
 
394
396
 
395
397
  Parameters
396
398
  ----------
397
- vars : Dict[str, str], default {}
398
- Dictionary of environment variables to set.
399
+ cpu : int, default 1
400
+ Number of CPUs required for this step.
401
+ gpu : int, optional, default None
402
+ Number of GPUs required for this step.
403
+ disk : int, optional, default None
404
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
405
+ memory : int, default 4096
406
+ Memory size (in MB) required for this step.
407
+ shared_memory : int, optional, default None
408
+ The value for the size (in MiB) of the /dev/shm volume for this step.
409
+ This parameter maps to the `--shm-size` option in Docker.
399
410
  """
400
411
  ...
401
412
 
402
413
  @typing.overload
403
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
414
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
404
415
  ...
405
416
 
406
417
  @typing.overload
407
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
418
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
408
419
  ...
409
420
 
410
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
421
+ 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):
411
422
  """
412
- Specifies environment variables to be set prior to the execution of a step.
423
+ Specifies the resources needed when executing this step.
413
424
 
425
+ Use `@resources` to specify the resource requirements
426
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
414
427
 
415
- Parameters
416
- ----------
417
- vars : Dict[str, str], default {}
418
- Dictionary of environment variables to set.
419
- """
420
- ...
421
-
422
- 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]]]:
423
- """
424
- Specifies that this step should execute on DGX cloud.
428
+ You can choose the compute layer on the command line by executing e.g.
429
+ ```
430
+ python myflow.py run --with batch
431
+ ```
432
+ or
433
+ ```
434
+ python myflow.py run --with kubernetes
435
+ ```
436
+ which executes the flow on the desired system using the
437
+ requirements specified in `@resources`.
425
438
 
426
439
 
427
440
  Parameters
428
441
  ----------
429
- gpu : int
430
- Number of GPUs to use.
431
- gpu_type : str
432
- Type of Nvidia GPU to use.
433
- queue_timeout : int
434
- Time to keep the job in NVCF's queue.
442
+ cpu : int, default 1
443
+ Number of CPUs required for this step.
444
+ gpu : int, optional, default None
445
+ Number of GPUs required for this step.
446
+ disk : int, optional, default None
447
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
448
+ memory : int, default 4096
449
+ Memory size (in MB) required for this step.
450
+ shared_memory : int, optional, default None
451
+ The value for the size (in MiB) of the /dev/shm volume for this step.
452
+ This parameter maps to the `--shm-size` option in Docker.
435
453
  """
436
454
  ...
437
455
 
438
456
  @typing.overload
439
- 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]]]:
457
+ 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]]]:
440
458
  """
441
- Enables checkpointing for a step.
442
-
459
+ Specifies environment variables to be set prior to the execution of a step.
443
460
 
444
461
 
445
462
  Parameters
446
463
  ----------
447
- load_policy : str, default: "fresh"
448
- The policy for loading the checkpoint. The following policies are supported:
449
- - "eager": Loads the the latest available checkpoint within the namespace.
450
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
451
- will be loaded at the start of the task.
452
- - "none": Do not load any checkpoint
453
- - "fresh": Loads the lastest checkpoint created within the running Task.
454
- This mode helps loading checkpoints across various retry attempts of the same task.
455
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
456
- created within the task will be loaded when the task is retries execution on failure.
457
-
458
- temp_dir_root : str, default: None
459
- The root directory under which `current.checkpoint.directory` will be created.
464
+ vars : Dict[str, str], default {}
465
+ Dictionary of environment variables to set.
460
466
  """
461
467
  ...
462
468
 
463
469
  @typing.overload
464
- def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
470
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
465
471
  ...
466
472
 
467
473
  @typing.overload
468
- def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
474
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
469
475
  ...
470
476
 
471
- 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):
477
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
472
478
  """
473
- Enables checkpointing for a step.
474
-
479
+ Specifies environment variables to be set prior to the execution of a step.
475
480
 
476
481
 
477
482
  Parameters
478
483
  ----------
479
- load_policy : str, default: "fresh"
480
- The policy for loading the checkpoint. The following policies are supported:
481
- - "eager": Loads the the latest available checkpoint within the namespace.
482
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
483
- will be loaded at the start of the task.
484
- - "none": Do not load any checkpoint
485
- - "fresh": Loads the lastest checkpoint created within the running Task.
486
- This mode helps loading checkpoints across various retry attempts of the same task.
487
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
488
- created within the task will be loaded when the task is retries execution on failure.
489
-
490
- temp_dir_root : str, default: None
491
- The root directory under which `current.checkpoint.directory` will be created.
492
- """
493
- ...
494
-
495
- @typing.overload
496
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
497
- """
498
- Decorator prototype for all step decorators. This function gets specialized
499
- and imported for all decorators types by _import_plugin_decorators().
500
- """
501
- ...
502
-
503
- @typing.overload
504
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
505
- ...
506
-
507
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
508
- """
509
- Decorator prototype for all step decorators. This function gets specialized
510
- and imported for all decorators types by _import_plugin_decorators().
484
+ vars : Dict[str, str], default {}
485
+ Dictionary of environment variables to set.
511
486
  """
512
487
  ...
513
488
 
514
- @typing.overload
515
- 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]]]:
489
+ def vllm(*, model: str, backend: str, openai_api_server: bool, debug: bool, card_refresh_interval: int, max_retries: int, retry_alert_frequency: int, engine_args: 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]]]:
516
490
  """
517
- Specifies the Conda environment for the step.
491
+ This decorator is used to run vllm APIs as Metaflow task sidecars.
518
492
 
519
- Information in this decorator will augment any
520
- attributes set in the `@conda_base` flow-level decorator. Hence,
521
- you can use `@conda_base` to set packages required by all
522
- steps and use `@conda` to specify step-specific overrides.
493
+ User code call
494
+ --------------
495
+ @vllm(
496
+ model="...",
497
+ ...
498
+ )
523
499
 
500
+ Valid backend options
501
+ ---------------------
502
+ - 'local': Run as a separate process on the local task machine.
524
503
 
525
- Parameters
526
- ----------
527
- packages : Dict[str, str], default {}
528
- Packages to use for this step. The key is the name of the package
529
- and the value is the version to use.
530
- libraries : Dict[str, str], default {}
531
- Supported for backward compatibility. When used with packages, packages will take precedence.
532
- python : str, optional, default None
533
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
534
- that the version used will correspond to the version of the Python interpreter used to start the run.
535
- disabled : bool, default False
536
- If set to True, disables @conda.
537
- """
538
- ...
539
-
540
- @typing.overload
541
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
542
- ...
543
-
544
- @typing.overload
545
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
546
- ...
547
-
548
- 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):
549
- """
550
- Specifies the Conda environment for the step.
504
+ Valid model options
505
+ -------------------
506
+ Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
551
507
 
552
- Information in this decorator will augment any
553
- attributes set in the `@conda_base` flow-level decorator. Hence,
554
- you can use `@conda_base` to set packages required by all
555
- steps and use `@conda` to specify step-specific overrides.
508
+ NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
509
+ If you need multiple models, you must create multiple @vllm decorators.
556
510
 
557
511
 
558
512
  Parameters
559
513
  ----------
560
- packages : Dict[str, str], default {}
561
- Packages to use for this step. The key is the name of the package
562
- and the value is the version to use.
563
- libraries : Dict[str, str], default {}
564
- Supported for backward compatibility. When used with packages, packages will take precedence.
565
- python : str, optional, default None
566
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
567
- that the version used will correspond to the version of the Python interpreter used to start the run.
568
- disabled : bool, default False
569
- If set to True, disables @conda.
514
+ model: str
515
+ HuggingFace model identifier to be served by vLLM.
516
+ backend: str
517
+ Determines where and how to run the vLLM process.
518
+ openai_api_server: bool
519
+ Whether to use OpenAI-compatible API server mode (subprocess) instead of native engine.
520
+ Default is False (uses native engine).
521
+ Set to True for backward compatibility with existing code.
522
+ debug: bool
523
+ Whether to turn on verbose debugging logs.
524
+ card_refresh_interval: int
525
+ Interval in seconds for refreshing the vLLM status card.
526
+ Only used when openai_api_server=True.
527
+ max_retries: int
528
+ Maximum number of retries checking for vLLM server startup.
529
+ Only used when openai_api_server=True.
530
+ retry_alert_frequency: int
531
+ Frequency of alert logs for vLLM server startup retries.
532
+ Only used when openai_api_server=True.
533
+ engine_args : dict
534
+ Additional keyword arguments to pass to the vLLM engine.
535
+ For example, `tensor_parallel_size=2`.
570
536
  """
571
537
  ...
572
538
 
573
539
  @typing.overload
574
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
540
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
575
541
  """
576
- Internal decorator to support Fast bakery
542
+ Decorator prototype for all step decorators. This function gets specialized
543
+ and imported for all decorators types by _import_plugin_decorators().
577
544
  """
578
545
  ...
579
546
 
580
547
  @typing.overload
581
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
548
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
582
549
  ...
583
550
 
584
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
551
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
585
552
  """
586
- Internal decorator to support Fast bakery
553
+ Decorator prototype for all step decorators. This function gets specialized
554
+ and imported for all decorators types by _import_plugin_decorators().
587
555
  """
588
556
  ...
589
557
 
590
- 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]]]:
558
+ 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]]]:
591
559
  """
592
- This decorator is used to run vllm APIs as Metaflow task sidecars.
560
+ This decorator is used to run Ollama APIs as Metaflow task sidecars.
593
561
 
594
562
  User code call
595
563
  --------------
596
- @vllm(
597
- model="...",
564
+ @ollama(
565
+ models=[...],
598
566
  ...
599
567
  )
600
568
 
601
569
  Valid backend options
602
570
  ---------------------
603
571
  - 'local': Run as a separate process on the local task machine.
572
+ - (TODO) 'managed': Outerbounds hosts and selects compute provider.
573
+ - (TODO) 'remote': Spin up separate instance to serve Ollama models.
604
574
 
605
575
  Valid model options
606
576
  -------------------
607
- Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
608
-
609
- NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
610
- If you need multiple models, you must create multiple @vllm decorators.
577
+ Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
611
578
 
612
579
 
613
580
  Parameters
614
581
  ----------
615
- model: str
616
- HuggingFace model identifier to be served by vLLM.
582
+ models: list[str]
583
+ List of Ollama containers running models in sidecars.
617
584
  backend: str
618
- Determines where and how to run the vLLM process.
585
+ Determines where and how to run the Ollama process.
586
+ force_pull: bool
587
+ Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
588
+ cache_update_policy: str
589
+ Cache update policy: "auto", "force", or "never".
590
+ force_cache_update: bool
591
+ Simple override for "force" cache update policy.
619
592
  debug: bool
620
593
  Whether to turn on verbose debugging logs.
621
- kwargs : Any
622
- Any other keyword arguments are passed directly to the vLLM engine.
623
- This allows for flexible configuration of vLLM server settings.
624
- For example, `tensor_parallel_size=2`.
594
+ circuit_breaker_config: dict
595
+ Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
596
+ timeout_config: dict
597
+ Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
625
598
  """
626
599
  ...
627
600
 
628
601
  @typing.overload
629
- def timeout(*, seconds: int = 0, minutes: int = 0, hours: int = 0) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
602
+ 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]]]:
630
603
  """
631
- Specifies a timeout for your step.
604
+ Specifies secrets to be retrieved and injected as environment variables prior to
605
+ the execution of a step.
606
+
607
+
608
+ Parameters
609
+ ----------
610
+ sources : List[Union[str, Dict[str, Any]]], default: []
611
+ List of secret specs, defining how the secrets are to be retrieved
612
+ """
613
+ ...
614
+
615
+ @typing.overload
616
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
617
+ ...
618
+
619
+ @typing.overload
620
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
621
+ ...
622
+
623
+ 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]]] = []):
624
+ """
625
+ Specifies secrets to be retrieved and injected as environment variables prior to
626
+ the execution of a step.
627
+
628
+
629
+ Parameters
630
+ ----------
631
+ sources : List[Union[str, Dict[str, Any]]], default: []
632
+ List of secret specs, defining how the secrets are to be retrieved
633
+ """
634
+ ...
635
+
636
+ @typing.overload
637
+ def timeout(*, seconds: int = 0, minutes: int = 0, hours: int = 0) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
638
+ """
639
+ Specifies a timeout for your step.
632
640
 
633
641
  This decorator is useful if this step may hang indefinitely.
634
642
 
@@ -684,152 +692,167 @@ def timeout(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None],
684
692
  """
685
693
  ...
686
694
 
687
- 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]]]:
695
+ @typing.overload
696
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
688
697
  """
689
- Decorator that helps cache, version and store models/datasets from huggingface hub.
690
-
691
-
692
- Parameters
693
- ----------
694
- temp_dir_root : str, optional
695
- The root directory that will hold the temporary directory where objects will be downloaded.
696
-
697
- load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
698
- The list of repos (models/datasets) to load.
699
-
700
- Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
701
-
702
- - If repo (model/dataset) is not found in the datastore:
703
- - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
704
- - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
705
- - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
706
-
707
- - If repo is found in the datastore:
708
- - Loads it directly from datastore to local path (can be temporary directory or specified path)
698
+ Internal decorator to support Fast bakery
709
699
  """
710
700
  ...
711
701
 
712
702
  @typing.overload
713
- 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]]]:
703
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
704
+ ...
705
+
706
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
714
707
  """
715
- Specifies secrets to be retrieved and injected as environment variables prior to
716
- the execution of a step.
708
+ Internal decorator to support Fast bakery
709
+ """
710
+ ...
711
+
712
+ @typing.overload
713
+ 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]]]:
714
+ """
715
+ Enables checkpointing for a step.
716
+
717
+ > Examples
718
+
719
+ - Saving Checkpoints
720
+
721
+ ```python
722
+ @checkpoint
723
+ @step
724
+ def train(self):
725
+ model = create_model(self.parameters, checkpoint_path = None)
726
+ for i in range(self.epochs):
727
+ # some training logic
728
+ loss = model.train(self.dataset)
729
+ if i % 10 == 0:
730
+ model.save(
731
+ current.checkpoint.directory,
732
+ )
733
+ # saves the contents of the `current.checkpoint.directory` as a checkpoint
734
+ # and returns a reference dictionary to the checkpoint saved in the datastore
735
+ self.latest_checkpoint = current.checkpoint.save(
736
+ name="epoch_checkpoint",
737
+ metadata={
738
+ "epoch": i,
739
+ "loss": loss,
740
+ }
741
+ )
742
+ ```
743
+
744
+ - Using Loaded Checkpoints
745
+
746
+ ```python
747
+ @retry(times=3)
748
+ @checkpoint
749
+ @step
750
+ def train(self):
751
+ # Assume that the task has restarted and the previous attempt of the task
752
+ # saved a checkpoint
753
+ checkpoint_path = None
754
+ if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
755
+ print("Loaded checkpoint from the previous attempt")
756
+ checkpoint_path = current.checkpoint.directory
757
+
758
+ model = create_model(self.parameters, checkpoint_path = checkpoint_path)
759
+ for i in range(self.epochs):
760
+ ...
761
+ ```
717
762
 
718
763
 
719
764
  Parameters
720
765
  ----------
721
- sources : List[Union[str, Dict[str, Any]]], default: []
722
- List of secret specs, defining how the secrets are to be retrieved
766
+ load_policy : str, default: "fresh"
767
+ The policy for loading the checkpoint. The following policies are supported:
768
+ - "eager": Loads the the latest available checkpoint within the namespace.
769
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
770
+ will be loaded at the start of the task.
771
+ - "none": Do not load any checkpoint
772
+ - "fresh": Loads the lastest checkpoint created within the running Task.
773
+ This mode helps loading checkpoints across various retry attempts of the same task.
774
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
775
+ created within the task will be loaded when the task is retries execution on failure.
776
+
777
+ temp_dir_root : str, default: None
778
+ The root directory under which `current.checkpoint.directory` will be created.
723
779
  """
724
780
  ...
725
781
 
726
782
  @typing.overload
727
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
783
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
728
784
  ...
729
785
 
730
786
  @typing.overload
731
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
787
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
732
788
  ...
733
789
 
734
- 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]]] = []):
790
+ 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):
735
791
  """
736
- Specifies secrets to be retrieved and injected as environment variables prior to
737
- the execution of a step.
792
+ Enables checkpointing for a step.
738
793
 
794
+ > Examples
739
795
 
740
- Parameters
741
- ----------
742
- sources : List[Union[str, Dict[str, Any]]], default: []
743
- List of secret specs, defining how the secrets are to be retrieved
744
- """
745
- ...
746
-
747
- def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: typing.Optional[str] = None, image_pull_policy: str = 'KUBERNETES_IMAGE_PULL_POLICY', image_pull_secrets: typing.List[str] = [], service_account: str = 'METAFLOW_KUBERNETES_SERVICE_ACCOUNT', secrets: typing.Optional[typing.List[str]] = None, node_selector: typing.Union[typing.Dict[str, str], str, None] = None, namespace: str = 'METAFLOW_KUBERNETES_NAMESPACE', gpu: typing.Optional[int] = None, gpu_vendor: str = 'KUBERNETES_GPU_VENDOR', tolerations: typing.List[str] = [], labels: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_LABELS', annotations: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_ANNOTATIONS', use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = '/metaflow_temp', persistent_volume_claims: typing.Optional[typing.Dict[str, str]] = None, shared_memory: typing.Optional[int] = None, port: typing.Optional[int] = None, compute_pool: typing.Optional[str] = None, hostname_resolution_timeout: int = 600, qos: str = 'Burstable', security_context: typing.Optional[typing.Dict[str, typing.Any]] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
748
- """
749
- Specifies that this step should execute on Kubernetes.
796
+ - Saving Checkpoints
797
+
798
+ ```python
799
+ @checkpoint
800
+ @step
801
+ def train(self):
802
+ model = create_model(self.parameters, checkpoint_path = None)
803
+ for i in range(self.epochs):
804
+ # some training logic
805
+ loss = model.train(self.dataset)
806
+ if i % 10 == 0:
807
+ model.save(
808
+ current.checkpoint.directory,
809
+ )
810
+ # saves the contents of the `current.checkpoint.directory` as a checkpoint
811
+ # and returns a reference dictionary to the checkpoint saved in the datastore
812
+ self.latest_checkpoint = current.checkpoint.save(
813
+ name="epoch_checkpoint",
814
+ metadata={
815
+ "epoch": i,
816
+ "loss": loss,
817
+ }
818
+ )
819
+ ```
820
+
821
+ - Using Loaded Checkpoints
822
+
823
+ ```python
824
+ @retry(times=3)
825
+ @checkpoint
826
+ @step
827
+ def train(self):
828
+ # Assume that the task has restarted and the previous attempt of the task
829
+ # saved a checkpoint
830
+ checkpoint_path = None
831
+ if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
832
+ print("Loaded checkpoint from the previous attempt")
833
+ checkpoint_path = current.checkpoint.directory
834
+
835
+ model = create_model(self.parameters, checkpoint_path = checkpoint_path)
836
+ for i in range(self.epochs):
837
+ ...
838
+ ```
750
839
 
751
840
 
752
841
  Parameters
753
842
  ----------
754
- cpu : int, default 1
755
- Number of CPUs required for this step. If `@resources` is
756
- also present, the maximum value from all decorators is used.
757
- memory : int, default 4096
758
- Memory size (in MB) required for this step. If
759
- `@resources` is also present, the maximum value from all decorators is
760
- used.
761
- disk : int, default 10240
762
- Disk size (in MB) required for this step. If
763
- `@resources` is also present, the maximum value from all decorators is
764
- used.
765
- image : str, optional, default None
766
- Docker image to use when launching on Kubernetes. If not specified, and
767
- METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
768
- not, a default Docker image mapping to the current version of Python is used.
769
- image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
770
- If given, the imagePullPolicy to be applied to the Docker image of the step.
771
- image_pull_secrets: List[str], default []
772
- The default is extracted from METAFLOW_KUBERNETES_IMAGE_PULL_SECRETS.
773
- Kubernetes image pull secrets to use when pulling container images
774
- in Kubernetes.
775
- service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
776
- Kubernetes service account to use when launching pod in Kubernetes.
777
- secrets : List[str], optional, default None
778
- Kubernetes secrets to use when launching pod in Kubernetes. These
779
- secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
780
- in Metaflow configuration.
781
- node_selector: Union[Dict[str,str], str], optional, default None
782
- Kubernetes node selector(s) to apply to the pod running the task.
783
- Can be passed in as a comma separated string of values e.g.
784
- 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
785
- {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
786
- namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
787
- Kubernetes namespace to use when launching pod in Kubernetes.
788
- gpu : int, optional, default None
789
- Number of GPUs required for this step. A value of zero implies that
790
- the scheduled node should not have GPUs.
791
- gpu_vendor : str, default KUBERNETES_GPU_VENDOR
792
- The vendor of the GPUs to be used for this step.
793
- tolerations : List[str], default []
794
- The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
795
- Kubernetes tolerations to use when launching pod in Kubernetes.
796
- labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
797
- Kubernetes labels to use when launching pod in Kubernetes.
798
- annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
799
- Kubernetes annotations to use when launching pod in Kubernetes.
800
- use_tmpfs : bool, default False
801
- This enables an explicit tmpfs mount for this step.
802
- tmpfs_tempdir : bool, default True
803
- sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
804
- tmpfs_size : int, optional, default: None
805
- The value for the size (in MiB) of the tmpfs mount for this step.
806
- This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
807
- memory allocated for this step.
808
- tmpfs_path : str, optional, default /metaflow_temp
809
- Path to tmpfs mount for this step.
810
- persistent_volume_claims : Dict[str, str], optional, default None
811
- A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
812
- volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
813
- shared_memory: int, optional
814
- Shared memory size (in MiB) required for this step
815
- port: int, optional
816
- Port number to specify in the Kubernetes job object
817
- compute_pool : str, optional, default None
818
- Compute pool to be used for for this step.
819
- If not specified, any accessible compute pool within the perimeter is used.
820
- hostname_resolution_timeout: int, default 10 * 60
821
- Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
822
- Only applicable when @parallel is used.
823
- qos: str, default: Burstable
824
- Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
843
+ load_policy : str, default: "fresh"
844
+ The policy for loading the checkpoint. The following policies are supported:
845
+ - "eager": Loads the the latest available checkpoint within the namespace.
846
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
847
+ will be loaded at the start of the task.
848
+ - "none": Do not load any checkpoint
849
+ - "fresh": Loads the lastest checkpoint created within the running Task.
850
+ This mode helps loading checkpoints across various retry attempts of the same task.
851
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
852
+ created within the task will be loaded when the task is retries execution on failure.
825
853
 
826
- security_context: Dict[str, Any], optional, default None
827
- Container security context. Applies to the task container. Allows the following keys:
828
- - privileged: bool, optional, default None
829
- - allow_privilege_escalation: bool, optional, default None
830
- - run_as_user: int, optional, default None
831
- - run_as_group: int, optional, default None
832
- - run_as_non_root: bool, optional, default None
854
+ temp_dir_root : str, default: None
855
+ The root directory under which `current.checkpoint.directory` will be created.
833
856
  """
834
857
  ...
835
858
 
@@ -888,323 +911,337 @@ def retry(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
888
911
  """
889
912
  ...
890
913
 
891
- @typing.overload
892
- 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]]]:
914
+ 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]]]:
893
915
  """
894
- Creates a human-readable report, a Metaflow Card, after this step completes.
895
-
896
- Note that you may add multiple `@card` decorators in a step with different parameters.
916
+ Specifies that this step is used to deploy an instance of the app.
917
+ Requires that self.app_name, self.app_port, self.entrypoint and self.deployDir is set.
897
918
 
898
919
 
899
920
  Parameters
900
921
  ----------
901
- type : str, default 'default'
902
- Card type.
903
- id : str, optional, default None
904
- If multiple cards are present, use this id to identify this card.
905
- options : Dict[str, Any], default {}
906
- Options passed to the card. The contents depend on the card type.
907
- timeout : int, default 45
908
- Interrupt reporting if it takes more than this many seconds.
922
+ app_port : int
923
+ Number of GPUs to use.
924
+ app_name : str
925
+ Name of the app to deploy.
909
926
  """
910
927
  ...
911
928
 
912
929
  @typing.overload
913
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
914
- ...
915
-
916
- @typing.overload
917
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
918
- ...
919
-
920
- 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):
930
+ 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]]]:
921
931
  """
922
- Creates a human-readable report, a Metaflow Card, after this step completes.
932
+ Specifies the Conda environment for the step.
923
933
 
924
- Note that you may add multiple `@card` decorators in a step with different parameters.
934
+ Information in this decorator will augment any
935
+ attributes set in the `@conda_base` flow-level decorator. Hence,
936
+ you can use `@conda_base` to set packages required by all
937
+ steps and use `@conda` to specify step-specific overrides.
925
938
 
926
939
 
927
940
  Parameters
928
941
  ----------
929
- type : str, default 'default'
930
- Card type.
931
- id : str, optional, default None
932
- If multiple cards are present, use this id to identify this card.
933
- options : Dict[str, Any], default {}
934
- Options passed to the card. The contents depend on the card type.
935
- timeout : int, default 45
936
- Interrupt reporting if it takes more than this many seconds.
942
+ packages : Dict[str, str], default {}
943
+ Packages to use for this step. The key is the name of the package
944
+ and the value is the version to use.
945
+ libraries : Dict[str, str], default {}
946
+ Supported for backward compatibility. When used with packages, packages will take precedence.
947
+ python : str, optional, default None
948
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
949
+ that the version used will correspond to the version of the Python interpreter used to start the run.
950
+ disabled : bool, default False
951
+ If set to True, disables @conda.
937
952
  """
938
953
  ...
939
954
 
940
955
  @typing.overload
941
- 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]]]:
956
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
957
+ ...
958
+
959
+ @typing.overload
960
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
961
+ ...
962
+
963
+ 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):
942
964
  """
943
- Specifies the resources needed when executing this step.
965
+ Specifies the Conda environment for the step.
944
966
 
945
- Use `@resources` to specify the resource requirements
946
- independently of the specific compute layer (`@batch`, `@kubernetes`).
967
+ Information in this decorator will augment any
968
+ attributes set in the `@conda_base` flow-level decorator. Hence,
969
+ you can use `@conda_base` to set packages required by all
970
+ steps and use `@conda` to specify step-specific overrides.
947
971
 
948
- You can choose the compute layer on the command line by executing e.g.
949
- ```
950
- python myflow.py run --with batch
951
- ```
952
- or
953
- ```
954
- python myflow.py run --with kubernetes
955
- ```
956
- which executes the flow on the desired system using the
957
- requirements specified in `@resources`.
972
+
973
+ Parameters
974
+ ----------
975
+ packages : Dict[str, str], default {}
976
+ Packages to use for this step. The key is the name of the package
977
+ and the value is the version to use.
978
+ libraries : Dict[str, str], default {}
979
+ Supported for backward compatibility. When used with packages, packages will take precedence.
980
+ python : str, optional, default None
981
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
982
+ that the version used will correspond to the version of the Python interpreter used to start the run.
983
+ disabled : bool, default False
984
+ If set to True, disables @conda.
985
+ """
986
+ ...
987
+
988
+ @typing.overload
989
+ 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]]]:
990
+ """
991
+ Specifies the PyPI packages for the step.
992
+
993
+ Information in this decorator will augment any
994
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
995
+ you can use `@pypi_base` to set packages required by all
996
+ steps and use `@pypi` to specify step-specific overrides.
958
997
 
959
998
 
960
999
  Parameters
961
1000
  ----------
962
- cpu : int, default 1
963
- Number of CPUs required for this step.
964
- gpu : int, optional, default None
965
- Number of GPUs required for this step.
966
- disk : int, optional, default None
967
- Disk size (in MB) required for this step. Only applies on Kubernetes.
968
- memory : int, default 4096
969
- Memory size (in MB) required for this step.
970
- shared_memory : int, optional, default None
971
- The value for the size (in MiB) of the /dev/shm volume for this step.
972
- This parameter maps to the `--shm-size` option in Docker.
1001
+ packages : Dict[str, str], default: {}
1002
+ Packages to use for this step. The key is the name of the package
1003
+ and the value is the version to use.
1004
+ python : str, optional, default: None
1005
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1006
+ that the version used will correspond to the version of the Python interpreter used to start the run.
973
1007
  """
974
1008
  ...
975
1009
 
976
1010
  @typing.overload
977
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1011
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
978
1012
  ...
979
1013
 
980
1014
  @typing.overload
981
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1015
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
982
1016
  ...
983
1017
 
984
- 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):
1018
+ 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):
985
1019
  """
986
- Specifies the resources needed when executing this step.
987
-
988
- Use `@resources` to specify the resource requirements
989
- independently of the specific compute layer (`@batch`, `@kubernetes`).
1020
+ Specifies the PyPI packages for the step.
990
1021
 
991
- You can choose the compute layer on the command line by executing e.g.
992
- ```
993
- python myflow.py run --with batch
994
- ```
995
- or
996
- ```
997
- python myflow.py run --with kubernetes
998
- ```
999
- which executes the flow on the desired system using the
1000
- requirements specified in `@resources`.
1022
+ Information in this decorator will augment any
1023
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
1024
+ you can use `@pypi_base` to set packages required by all
1025
+ steps and use `@pypi` to specify step-specific overrides.
1001
1026
 
1002
1027
 
1003
1028
  Parameters
1004
1029
  ----------
1005
- cpu : int, default 1
1006
- Number of CPUs required for this step.
1007
- gpu : int, optional, default None
1008
- Number of GPUs required for this step.
1009
- disk : int, optional, default None
1010
- Disk size (in MB) required for this step. Only applies on Kubernetes.
1011
- memory : int, default 4096
1012
- Memory size (in MB) required for this step.
1013
- shared_memory : int, optional, default None
1014
- The value for the size (in MiB) of the /dev/shm volume for this step.
1015
- This parameter maps to the `--shm-size` option in Docker.
1030
+ packages : Dict[str, str], default: {}
1031
+ Packages to use for this step. The key is the name of the package
1032
+ and the value is the version to use.
1033
+ python : str, optional, default: None
1034
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1035
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1016
1036
  """
1017
1037
  ...
1018
1038
 
1019
- def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1039
+ @typing.overload
1040
+ 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]]]:
1020
1041
  """
1021
- Specifies what flows belong to the same project.
1042
+ Enables loading / saving of models within a step.
1022
1043
 
1023
- A project-specific namespace is created for all flows that
1024
- use the same `@project(name)`.
1044
+ > Examples
1045
+ - Saving Models
1046
+ ```python
1047
+ @model
1048
+ @step
1049
+ def train(self):
1050
+ # current.model.save returns a dictionary reference to the model saved
1051
+ self.my_model = current.model.save(
1052
+ path_to_my_model,
1053
+ label="my_model",
1054
+ metadata={
1055
+ "epochs": 10,
1056
+ "batch-size": 32,
1057
+ "learning-rate": 0.001,
1058
+ }
1059
+ )
1060
+ self.next(self.test)
1061
+
1062
+ @model(load="my_model")
1063
+ @step
1064
+ def test(self):
1065
+ # `current.model.loaded` returns a dictionary of the loaded models
1066
+ # where the key is the name of the artifact and the value is the path to the model
1067
+ print(os.listdir(current.model.loaded["my_model"]))
1068
+ self.next(self.end)
1069
+ ```
1070
+
1071
+ - Loading models
1072
+ ```python
1073
+ @step
1074
+ def train(self):
1075
+ # current.model.load returns the path to the model loaded
1076
+ checkpoint_path = current.model.load(
1077
+ self.checkpoint_key,
1078
+ )
1079
+ model_path = current.model.load(
1080
+ self.model,
1081
+ )
1082
+ self.next(self.test)
1083
+ ```
1025
1084
 
1026
1085
 
1027
1086
  Parameters
1028
1087
  ----------
1029
- name : str
1030
- Project name. Make sure that the name is unique amongst all
1031
- projects that use the same production scheduler. The name may
1032
- contain only lowercase alphanumeric characters and underscores.
1033
-
1034
- branch : Optional[str], default None
1035
- The branch to use. If not specified, the branch is set to
1036
- `user.<username>` unless `production` is set to `True`. This can
1037
- also be set on the command line using `--branch` as a top-level option.
1038
- It is an error to specify `branch` in the decorator and on the command line.
1088
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
1089
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
1090
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by `current.checkpoint` / `current.model` / `current.huggingface_hub`.
1091
+ 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
1092
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
1093
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
1039
1094
 
1040
- production : bool, default False
1041
- Whether or not the branch is the production branch. This can also be set on the
1042
- command line using `--production` as a top-level option. It is an error to specify
1043
- `production` in the decorator and on the command line.
1044
- The project branch name will be:
1045
- - if `branch` is specified:
1046
- - if `production` is True: `prod.<branch>`
1047
- - if `production` is False: `test.<branch>`
1048
- - if `branch` is not specified:
1049
- - if `production` is True: `prod`
1050
- - if `production` is False: `user.<username>`
1095
+ temp_dir_root : str, default: None
1096
+ The root directory under which `current.model.loaded` will store loaded models
1051
1097
  """
1052
1098
  ...
1053
1099
 
1054
- 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]]:
1100
+ @typing.overload
1101
+ def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1102
+ ...
1103
+
1104
+ @typing.overload
1105
+ def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1106
+ ...
1107
+
1108
+ 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):
1055
1109
  """
1056
- 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.
1057
- 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.
1110
+ Enables loading / saving of models within a step.
1111
+
1112
+ > Examples
1113
+ - Saving Models
1114
+ ```python
1115
+ @model
1116
+ @step
1117
+ def train(self):
1118
+ # current.model.save returns a dictionary reference to the model saved
1119
+ self.my_model = current.model.save(
1120
+ path_to_my_model,
1121
+ label="my_model",
1122
+ metadata={
1123
+ "epochs": 10,
1124
+ "batch-size": 32,
1125
+ "learning-rate": 0.001,
1126
+ }
1127
+ )
1128
+ self.next(self.test)
1129
+
1130
+ @model(load="my_model")
1131
+ @step
1132
+ def test(self):
1133
+ # `current.model.loaded` returns a dictionary of the loaded models
1134
+ # where the key is the name of the artifact and the value is the path to the model
1135
+ print(os.listdir(current.model.loaded["my_model"]))
1136
+ self.next(self.end)
1137
+ ```
1138
+
1139
+ - Loading models
1140
+ ```python
1141
+ @step
1142
+ def train(self):
1143
+ # current.model.load returns the path to the model loaded
1144
+ checkpoint_path = current.model.load(
1145
+ self.checkpoint_key,
1146
+ )
1147
+ model_path = current.model.load(
1148
+ self.model,
1149
+ )
1150
+ self.next(self.test)
1151
+ ```
1058
1152
 
1059
1153
 
1060
1154
  Parameters
1061
1155
  ----------
1062
- timeout : int
1063
- Time, in seconds before the task times out and fails. (Default: 3600)
1064
- poke_interval : int
1065
- Time in seconds that the job should wait in between each try. (Default: 60)
1066
- mode : str
1067
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1068
- exponential_backoff : bool
1069
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1070
- pool : str
1071
- the slot pool this task should run in,
1072
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1073
- soft_fail : bool
1074
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1075
- name : str
1076
- Name of the sensor on Airflow
1077
- description : str
1078
- Description of sensor in the Airflow UI
1079
- external_dag_id : str
1080
- The dag_id that contains the task you want to wait for.
1081
- external_task_ids : List[str]
1082
- The list of task_ids that you want to wait for.
1083
- If None (default value) the sensor waits for the DAG. (Default: None)
1084
- allowed_states : List[str]
1085
- Iterable of allowed states, (Default: ['success'])
1086
- failed_states : List[str]
1087
- Iterable of failed or dis-allowed states. (Default: None)
1088
- execution_delta : datetime.timedelta
1089
- time difference with the previous execution to look at,
1090
- the default is the same logical date as the current task or DAG. (Default: None)
1091
- check_existence: bool
1092
- Set to True to check if the external task exists or check if
1093
- the DAG to wait for exists. (Default: True)
1156
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
1157
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
1158
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by `current.checkpoint` / `current.model` / `current.huggingface_hub`.
1159
+ 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
1160
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
1161
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
1162
+
1163
+ temp_dir_root : str, default: None
1164
+ The root directory under which `current.model.loaded` will store loaded models
1094
1165
  """
1095
1166
  ...
1096
1167
 
1097
- def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None):
1168
+ 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]]]:
1098
1169
  """
1099
- Allows setting external datastores to save data for the
1100
- `@checkpoint`/`@model`/`@huggingface_hub` decorators.
1101
-
1102
- This decorator is useful when users wish to save data to a different datastore
1103
- than what is configured in Metaflow. This can be for variety of reasons:
1170
+ Decorator that helps cache, version and store models/datasets from huggingface hub.
1104
1171
 
1105
- 1. Data security: The objects needs to be stored in a bucket (object storage) that is not accessible by other flows.
1106
- 2. Data Locality: The location where the task is executing is not located in the same region as the datastore.
1107
- - Example: Metaflow datastore lives in US East, but the task is executing in Finland datacenters.
1108
- 3. Data Lifecycle Policies: The objects need to be archived / managed separately from the Metaflow managed objects.
1109
- - Example: Flow is training very large models that need to be stored separately and will be deleted more aggressively than the Metaflow managed objects.
1172
+ > Examples
1173
+
1174
+ **Usage: creating references of models from huggingface that may be loaded in downstream steps**
1175
+ ```python
1176
+ @huggingface_hub
1177
+ @step
1178
+ def pull_model_from_huggingface(self):
1179
+ # `current.huggingface_hub.snapshot_download` downloads the model from the Hugging Face Hub
1180
+ # and saves it in the backend storage based on the model's `repo_id`. If there exists a model
1181
+ # with the same `repo_id` in the backend storage, it will not download the model again. The return
1182
+ # value of the function is a reference to the model in the backend storage.
1183
+ # This reference can be used to load the model in the subsequent steps via `@model(load=["llama_model"])`
1184
+
1185
+ self.model_id = "mistralai/Mistral-7B-Instruct-v0.1"
1186
+ self.llama_model = current.huggingface_hub.snapshot_download(
1187
+ repo_id=self.model_id,
1188
+ allow_patterns=["*.safetensors", "*.json", "tokenizer.*"],
1189
+ )
1190
+ self.next(self.train)
1191
+ ```
1110
1192
 
1111
- Usage:
1112
- ----------
1193
+ **Usage: loading models directly from huggingface hub or from cache (from metaflow's datastore)**
1194
+ ```python
1195
+ @huggingface_hub(load=["mistralai/Mistral-7B-Instruct-v0.1"])
1196
+ @step
1197
+ def pull_model_from_huggingface(self):
1198
+ path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
1199
+ ```
1113
1200
 
1114
- - Using a custom IAM role to access the datastore.
1201
+ ```python
1202
+ @huggingface_hub(load=[("mistralai/Mistral-7B-Instruct-v0.1", "/my-directory"), ("myorg/mistral-lora, "/my-lora-directory")])
1203
+ @step
1204
+ def finetune_model(self):
1205
+ path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
1206
+ # path_to_model will be /my-directory
1207
+ ```
1115
1208
 
1116
- ```python
1117
- @with_artifact_store(
1118
- type="s3",
1119
- config=lambda: {
1120
- "root": "s3://my-bucket-foo/path/to/root",
1121
- "role_arn": ROLE,
1209
+ ```python
1210
+ # Takes all the arguments passed to `snapshot_download`
1211
+ # except for `local_dir`
1212
+ @huggingface_hub(load=[
1213
+ {
1214
+ "repo_id": "mistralai/Mistral-7B-Instruct-v0.1",
1122
1215
  },
1123
- )
1124
- class MyFlow(FlowSpec):
1125
-
1126
- @checkpoint
1127
- @step
1128
- def start(self):
1129
- with open("my_file.txt", "w") as f:
1130
- f.write("Hello, World!")
1131
- self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1132
- self.next(self.end)
1133
-
1134
- ```
1135
-
1136
- - Using credentials to access the s3-compatible datastore.
1137
-
1138
- ```python
1139
- @with_artifact_store(
1140
- type="s3",
1141
- config=lambda: {
1142
- "root": "s3://my-bucket-foo/path/to/root",
1143
- "client_params": {
1144
- "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1145
- "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1146
- },
1216
+ {
1217
+ "repo_id": "myorg/mistral-lora",
1218
+ "repo_type": "model",
1147
1219
  },
1148
- )
1149
- class MyFlow(FlowSpec):
1150
-
1151
- @checkpoint
1152
- @step
1153
- def start(self):
1154
- with open("my_file.txt", "w") as f:
1155
- f.write("Hello, World!")
1156
- self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1157
- self.next(self.end)
1220
+ ])
1221
+ @step
1222
+ def finetune_model(self):
1223
+ path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
1224
+ # path_to_model will be /my-directory
1225
+ ```
1158
1226
 
1159
- ```
1160
1227
 
1161
- - Accessing objects stored in external datastores after task execution.
1228
+ Parameters
1229
+ ----------
1230
+ temp_dir_root : str, optional
1231
+ The root directory that will hold the temporary directory where objects will be downloaded.
1162
1232
 
1163
- ```python
1164
- run = Run("CheckpointsTestsFlow/8992")
1165
- with artifact_store_from(run=run, config={
1166
- "client_params": {
1167
- "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1168
- "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1169
- },
1170
- }):
1171
- with Checkpoint() as cp:
1172
- latest = cp.list(
1173
- task=run["start"].task
1174
- )[0]
1175
- print(latest)
1176
- cp.load(
1177
- latest,
1178
- "test-checkpoints"
1179
- )
1233
+ load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
1234
+ The list of repos (models/datasets) to load.
1180
1235
 
1181
- task = Task("TorchTuneFlow/8484/train/53673")
1182
- with artifact_store_from(run=run, config={
1183
- "client_params": {
1184
- "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1185
- "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1186
- },
1187
- }):
1188
- load_model(
1189
- task.data.model_ref,
1190
- "test-models"
1191
- )
1192
- ```
1193
- Parameters:
1194
- ----------
1236
+ Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
1195
1237
 
1196
- type: str
1197
- The type of the datastore. Can be one of 's3', 'gcs', 'azure' or any other supported metaflow Datastore.
1238
+ - If repo (model/dataset) is not found in the datastore:
1239
+ - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
1240
+ - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
1241
+ - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
1198
1242
 
1199
- config: dict or Callable
1200
- Dictionary of configuration options for the datastore. The following keys are required:
1201
- - root: The root path in the datastore where the data will be saved. (needs to be in the format expected by the datastore)
1202
- - example: 's3://bucket-name/path/to/root'
1203
- - example: 'gs://bucket-name/path/to/root'
1204
- - example: 'https://myblockacc.blob.core.windows.net/metaflow/'
1205
- - role_arn (optional): AWS IAM role to access s3 bucket (only when `type` is 's3')
1206
- - session_vars (optional): AWS session variables to access s3 bucket (only when `type` is 's3')
1207
- - client_params (optional): AWS client parameters to access s3 bucket (only when `type` is 's3')
1243
+ - If repo is found in the datastore:
1244
+ - Loads it directly from datastore to local path (can be temporary directory or specified path)
1208
1245
  """
1209
1246
  ...
1210
1247
 
@@ -1250,104 +1287,247 @@ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packag
1250
1287
  ...
1251
1288
 
1252
1289
  @typing.overload
1253
- 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]]:
1290
+ 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]]:
1254
1291
  """
1255
- Specifies the times when the flow should be run when running on a
1256
- production scheduler.
1292
+ Specifies the event(s) that this flow depends on.
1293
+
1294
+ ```
1295
+ @trigger(event='foo')
1296
+ ```
1297
+ or
1298
+ ```
1299
+ @trigger(events=['foo', 'bar'])
1300
+ ```
1301
+
1302
+ Additionally, you can specify the parameter mappings
1303
+ to map event payload to Metaflow parameters for the flow.
1304
+ ```
1305
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1306
+ ```
1307
+ or
1308
+ ```
1309
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1310
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1311
+ ```
1312
+
1313
+ 'parameters' can also be a list of strings and tuples like so:
1314
+ ```
1315
+ @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1316
+ ```
1317
+ This is equivalent to:
1318
+ ```
1319
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1320
+ ```
1257
1321
 
1258
1322
 
1259
1323
  Parameters
1260
1324
  ----------
1261
- hourly : bool, default False
1262
- Run the workflow hourly.
1263
- daily : bool, default True
1264
- Run the workflow daily.
1265
- weekly : bool, default False
1266
- Run the workflow weekly.
1267
- cron : str, optional, default None
1268
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1269
- specified by this expression.
1270
- timezone : str, optional, default None
1271
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1272
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1325
+ event : Union[str, Dict[str, Any]], optional, default None
1326
+ Event dependency for this flow.
1327
+ events : List[Union[str, Dict[str, Any]]], default []
1328
+ Events dependency for this flow.
1329
+ options : Dict[str, Any], default {}
1330
+ Backend-specific configuration for tuning eventing behavior.
1273
1331
  """
1274
1332
  ...
1275
1333
 
1276
1334
  @typing.overload
1277
- def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1335
+ def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1278
1336
  ...
1279
1337
 
1280
- def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None):
1338
+ 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] = {}):
1281
1339
  """
1282
- Specifies the times when the flow should be run when running on a
1283
- production scheduler.
1340
+ Specifies the event(s) that this flow depends on.
1341
+
1342
+ ```
1343
+ @trigger(event='foo')
1344
+ ```
1345
+ or
1346
+ ```
1347
+ @trigger(events=['foo', 'bar'])
1348
+ ```
1349
+
1350
+ Additionally, you can specify the parameter mappings
1351
+ to map event payload to Metaflow parameters for the flow.
1352
+ ```
1353
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1354
+ ```
1355
+ or
1356
+ ```
1357
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1358
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1359
+ ```
1360
+
1361
+ 'parameters' can also be a list of strings and tuples like so:
1362
+ ```
1363
+ @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1364
+ ```
1365
+ This is equivalent to:
1366
+ ```
1367
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1368
+ ```
1284
1369
 
1285
1370
 
1286
1371
  Parameters
1287
1372
  ----------
1288
- hourly : bool, default False
1289
- Run the workflow hourly.
1290
- daily : bool, default True
1291
- Run the workflow daily.
1292
- weekly : bool, default False
1293
- Run the workflow weekly.
1294
- cron : str, optional, default None
1295
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1296
- specified by this expression.
1297
- timezone : str, optional, default None
1298
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1299
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1373
+ event : Union[str, Dict[str, Any]], optional, default None
1374
+ Event dependency for this flow.
1375
+ events : List[Union[str, Dict[str, Any]]], default []
1376
+ Events dependency for this flow.
1377
+ options : Dict[str, Any], default {}
1378
+ Backend-specific configuration for tuning eventing behavior.
1300
1379
  """
1301
1380
  ...
1302
1381
 
1303
1382
  @typing.overload
1304
- def conda_base(*, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1383
+ 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]]:
1305
1384
  """
1306
- Specifies the Conda environment for all steps of the flow.
1385
+ Specifies the flow(s) that this flow depends on.
1307
1386
 
1308
- Use `@conda_base` to set common libraries required by all
1309
- steps and use `@conda` to specify step-specific additions.
1387
+ ```
1388
+ @trigger_on_finish(flow='FooFlow')
1389
+ ```
1390
+ or
1391
+ ```
1392
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1393
+ ```
1394
+ This decorator respects the @project decorator and triggers the flow
1395
+ when upstream runs within the same namespace complete successfully
1396
+
1397
+ Additionally, you can specify project aware upstream flow dependencies
1398
+ by specifying the fully qualified project_flow_name.
1399
+ ```
1400
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1401
+ ```
1402
+ or
1403
+ ```
1404
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1405
+ ```
1406
+
1407
+ You can also specify just the project or project branch (other values will be
1408
+ inferred from the current project or project branch):
1409
+ ```
1410
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1411
+ ```
1412
+
1413
+ Note that `branch` is typically one of:
1414
+ - `prod`
1415
+ - `user.bob`
1416
+ - `test.my_experiment`
1417
+ - `prod.staging`
1310
1418
 
1311
1419
 
1312
1420
  Parameters
1313
1421
  ----------
1314
- packages : Dict[str, str], default {}
1315
- Packages to use for this flow. The key is the name of the package
1316
- and the value is the version to use.
1317
- libraries : Dict[str, str], default {}
1318
- Supported for backward compatibility. When used with packages, packages will take precedence.
1319
- python : str, optional, default None
1320
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1321
- that the version used will correspond to the version of the Python interpreter used to start the run.
1322
- disabled : bool, default False
1323
- If set to True, disables Conda.
1422
+ flow : Union[str, Dict[str, str]], optional, default None
1423
+ Upstream flow dependency for this flow.
1424
+ flows : List[Union[str, Dict[str, str]]], default []
1425
+ Upstream flow dependencies for this flow.
1426
+ options : Dict[str, Any], default {}
1427
+ Backend-specific configuration for tuning eventing behavior.
1324
1428
  """
1325
1429
  ...
1326
1430
 
1327
1431
  @typing.overload
1328
- def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1432
+ def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1329
1433
  ...
1330
1434
 
1331
- def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
1435
+ 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] = {}):
1332
1436
  """
1333
- Specifies the Conda environment for all steps of the flow.
1437
+ Specifies the flow(s) that this flow depends on.
1334
1438
 
1335
- Use `@conda_base` to set common libraries required by all
1336
- steps and use `@conda` to specify step-specific additions.
1439
+ ```
1440
+ @trigger_on_finish(flow='FooFlow')
1441
+ ```
1442
+ or
1443
+ ```
1444
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1445
+ ```
1446
+ This decorator respects the @project decorator and triggers the flow
1447
+ when upstream runs within the same namespace complete successfully
1448
+
1449
+ Additionally, you can specify project aware upstream flow dependencies
1450
+ by specifying the fully qualified project_flow_name.
1451
+ ```
1452
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1453
+ ```
1454
+ or
1455
+ ```
1456
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1457
+ ```
1458
+
1459
+ You can also specify just the project or project branch (other values will be
1460
+ inferred from the current project or project branch):
1461
+ ```
1462
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1463
+ ```
1464
+
1465
+ Note that `branch` is typically one of:
1466
+ - `prod`
1467
+ - `user.bob`
1468
+ - `test.my_experiment`
1469
+ - `prod.staging`
1337
1470
 
1338
1471
 
1339
1472
  Parameters
1340
1473
  ----------
1341
- packages : Dict[str, str], default {}
1342
- Packages to use for this flow. The key is the name of the package
1343
- and the value is the version to use.
1344
- libraries : Dict[str, str], default {}
1345
- Supported for backward compatibility. When used with packages, packages will take precedence.
1346
- python : str, optional, default None
1347
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1348
- that the version used will correspond to the version of the Python interpreter used to start the run.
1349
- disabled : bool, default False
1350
- If set to True, disables Conda.
1474
+ flow : Union[str, Dict[str, str]], optional, default None
1475
+ Upstream flow dependency for this flow.
1476
+ flows : List[Union[str, Dict[str, str]]], default []
1477
+ Upstream flow dependencies for this flow.
1478
+ options : Dict[str, Any], default {}
1479
+ Backend-specific configuration for tuning eventing behavior.
1480
+ """
1481
+ ...
1482
+
1483
+ @typing.overload
1484
+ 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]]:
1485
+ """
1486
+ Specifies the times when the flow should be run when running on a
1487
+ production scheduler.
1488
+
1489
+
1490
+ Parameters
1491
+ ----------
1492
+ hourly : bool, default False
1493
+ Run the workflow hourly.
1494
+ daily : bool, default True
1495
+ Run the workflow daily.
1496
+ weekly : bool, default False
1497
+ Run the workflow weekly.
1498
+ cron : str, optional, default None
1499
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1500
+ specified by this expression.
1501
+ timezone : str, optional, default None
1502
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1503
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1504
+ """
1505
+ ...
1506
+
1507
+ @typing.overload
1508
+ def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1509
+ ...
1510
+
1511
+ def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None):
1512
+ """
1513
+ Specifies the times when the flow should be run when running on a
1514
+ production scheduler.
1515
+
1516
+
1517
+ Parameters
1518
+ ----------
1519
+ hourly : bool, default False
1520
+ Run the workflow hourly.
1521
+ daily : bool, default True
1522
+ Run the workflow daily.
1523
+ weekly : bool, default False
1524
+ Run the workflow weekly.
1525
+ cron : str, optional, default None
1526
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1527
+ specified by this expression.
1528
+ timezone : str, optional, default None
1529
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1530
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1351
1531
  """
1352
1532
  ...
1353
1533
 
@@ -1395,196 +1575,245 @@ def airflow_s3_key_sensor(*, timeout: int, poke_interval: int, mode: str, expone
1395
1575
  ...
1396
1576
 
1397
1577
  @typing.overload
1398
- 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]]:
1578
+ def conda_base(*, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1399
1579
  """
1400
- Specifies the flow(s) that this flow depends on.
1401
-
1402
- ```
1403
- @trigger_on_finish(flow='FooFlow')
1404
- ```
1405
- or
1406
- ```
1407
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1408
- ```
1409
- This decorator respects the @project decorator and triggers the flow
1410
- when upstream runs within the same namespace complete successfully
1411
-
1412
- Additionally, you can specify project aware upstream flow dependencies
1413
- by specifying the fully qualified project_flow_name.
1414
- ```
1415
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1416
- ```
1417
- or
1418
- ```
1419
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1420
- ```
1421
-
1422
- You can also specify just the project or project branch (other values will be
1423
- inferred from the current project or project branch):
1424
- ```
1425
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1426
- ```
1580
+ Specifies the Conda environment for all steps of the flow.
1427
1581
 
1428
- Note that `branch` is typically one of:
1429
- - `prod`
1430
- - `user.bob`
1431
- - `test.my_experiment`
1432
- - `prod.staging`
1582
+ Use `@conda_base` to set common libraries required by all
1583
+ steps and use `@conda` to specify step-specific additions.
1433
1584
 
1434
1585
 
1435
1586
  Parameters
1436
1587
  ----------
1437
- flow : Union[str, Dict[str, str]], optional, default None
1438
- Upstream flow dependency for this flow.
1439
- flows : List[Union[str, Dict[str, str]]], default []
1440
- Upstream flow dependencies for this flow.
1441
- options : Dict[str, Any], default {}
1442
- Backend-specific configuration for tuning eventing behavior.
1588
+ packages : Dict[str, str], default {}
1589
+ Packages to use for this flow. The key is the name of the package
1590
+ and the value is the version to use.
1591
+ libraries : Dict[str, str], default {}
1592
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1593
+ python : str, optional, default None
1594
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1595
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1596
+ disabled : bool, default False
1597
+ If set to True, disables Conda.
1443
1598
  """
1444
1599
  ...
1445
1600
 
1446
1601
  @typing.overload
1447
- def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1602
+ def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1448
1603
  ...
1449
1604
 
1450
- 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] = {}):
1605
+ def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
1451
1606
  """
1452
- Specifies the flow(s) that this flow depends on.
1453
-
1454
- ```
1455
- @trigger_on_finish(flow='FooFlow')
1456
- ```
1457
- or
1458
- ```
1459
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1460
- ```
1461
- This decorator respects the @project decorator and triggers the flow
1462
- when upstream runs within the same namespace complete successfully
1607
+ Specifies the Conda environment for all steps of the flow.
1463
1608
 
1464
- Additionally, you can specify project aware upstream flow dependencies
1465
- by specifying the fully qualified project_flow_name.
1466
- ```
1467
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1468
- ```
1469
- or
1470
- ```
1471
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1472
- ```
1609
+ Use `@conda_base` to set common libraries required by all
1610
+ steps and use `@conda` to specify step-specific additions.
1473
1611
 
1474
- You can also specify just the project or project branch (other values will be
1475
- inferred from the current project or project branch):
1476
- ```
1477
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1478
- ```
1479
1612
 
1480
- Note that `branch` is typically one of:
1481
- - `prod`
1482
- - `user.bob`
1483
- - `test.my_experiment`
1484
- - `prod.staging`
1613
+ Parameters
1614
+ ----------
1615
+ packages : Dict[str, str], default {}
1616
+ Packages to use for this flow. The key is the name of the package
1617
+ and the value is the version to use.
1618
+ libraries : Dict[str, str], default {}
1619
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1620
+ python : str, optional, default None
1621
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1622
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1623
+ disabled : bool, default False
1624
+ If set to True, disables Conda.
1625
+ """
1626
+ ...
1627
+
1628
+ 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]]:
1629
+ """
1630
+ 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.
1631
+ 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.
1485
1632
 
1486
1633
 
1487
1634
  Parameters
1488
1635
  ----------
1489
- flow : Union[str, Dict[str, str]], optional, default None
1490
- Upstream flow dependency for this flow.
1491
- flows : List[Union[str, Dict[str, str]]], default []
1492
- Upstream flow dependencies for this flow.
1493
- options : Dict[str, Any], default {}
1494
- Backend-specific configuration for tuning eventing behavior.
1636
+ timeout : int
1637
+ Time, in seconds before the task times out and fails. (Default: 3600)
1638
+ poke_interval : int
1639
+ Time in seconds that the job should wait in between each try. (Default: 60)
1640
+ mode : str
1641
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1642
+ exponential_backoff : bool
1643
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1644
+ pool : str
1645
+ the slot pool this task should run in,
1646
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1647
+ soft_fail : bool
1648
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1649
+ name : str
1650
+ Name of the sensor on Airflow
1651
+ description : str
1652
+ Description of sensor in the Airflow UI
1653
+ external_dag_id : str
1654
+ The dag_id that contains the task you want to wait for.
1655
+ external_task_ids : List[str]
1656
+ The list of task_ids that you want to wait for.
1657
+ If None (default value) the sensor waits for the DAG. (Default: None)
1658
+ allowed_states : List[str]
1659
+ Iterable of allowed states, (Default: ['success'])
1660
+ failed_states : List[str]
1661
+ Iterable of failed or dis-allowed states. (Default: None)
1662
+ execution_delta : datetime.timedelta
1663
+ time difference with the previous execution to look at,
1664
+ the default is the same logical date as the current task or DAG. (Default: None)
1665
+ check_existence: bool
1666
+ Set to True to check if the external task exists or check if
1667
+ the DAG to wait for exists. (Default: True)
1495
1668
  """
1496
1669
  ...
1497
1670
 
1498
- @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]]:
1671
+ def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None):
1500
1672
  """
1501
- Specifies the event(s) that this flow depends on.
1673
+ Allows setting external datastores to save data for the
1674
+ `@checkpoint`/`@model`/`@huggingface_hub` decorators.
1502
1675
 
1503
- ```
1504
- @trigger(event='foo')
1505
- ```
1506
- or
1507
- ```
1508
- @trigger(events=['foo', 'bar'])
1509
- ```
1676
+ This decorator is useful when users wish to save data to a different datastore
1677
+ than what is configured in Metaflow. This can be for variety of reasons:
1510
1678
 
1511
- Additionally, you can specify the parameter mappings
1512
- to map event payload to Metaflow parameters for the flow.
1513
- ```
1514
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1515
- ```
1516
- or
1517
- ```
1518
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1519
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1520
- ```
1679
+ 1. Data security: The objects needs to be stored in a bucket (object storage) that is not accessible by other flows.
1680
+ 2. Data Locality: The location where the task is executing is not located in the same region as the datastore.
1681
+ - Example: Metaflow datastore lives in US East, but the task is executing in Finland datacenters.
1682
+ 3. Data Lifecycle Policies: The objects need to be archived / managed separately from the Metaflow managed objects.
1683
+ - Example: Flow is training very large models that need to be stored separately and will be deleted more aggressively than the Metaflow managed objects.
1521
1684
 
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:
1527
- ```
1528
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1529
- ```
1685
+ Usage:
1686
+ ----------
1530
1687
 
1688
+ - Using a custom IAM role to access the datastore.
1531
1689
 
1532
- Parameters
1690
+ ```python
1691
+ @with_artifact_store(
1692
+ type="s3",
1693
+ config=lambda: {
1694
+ "root": "s3://my-bucket-foo/path/to/root",
1695
+ "role_arn": ROLE,
1696
+ },
1697
+ )
1698
+ class MyFlow(FlowSpec):
1699
+
1700
+ @checkpoint
1701
+ @step
1702
+ def start(self):
1703
+ with open("my_file.txt", "w") as f:
1704
+ f.write("Hello, World!")
1705
+ self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1706
+ self.next(self.end)
1707
+
1708
+ ```
1709
+
1710
+ - Using credentials to access the s3-compatible datastore.
1711
+
1712
+ ```python
1713
+ @with_artifact_store(
1714
+ type="s3",
1715
+ config=lambda: {
1716
+ "root": "s3://my-bucket-foo/path/to/root",
1717
+ "client_params": {
1718
+ "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1719
+ "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1720
+ },
1721
+ },
1722
+ )
1723
+ class MyFlow(FlowSpec):
1724
+
1725
+ @checkpoint
1726
+ @step
1727
+ def start(self):
1728
+ with open("my_file.txt", "w") as f:
1729
+ f.write("Hello, World!")
1730
+ self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1731
+ self.next(self.end)
1732
+
1733
+ ```
1734
+
1735
+ - Accessing objects stored in external datastores after task execution.
1736
+
1737
+ ```python
1738
+ run = Run("CheckpointsTestsFlow/8992")
1739
+ with artifact_store_from(run=run, config={
1740
+ "client_params": {
1741
+ "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1742
+ "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1743
+ },
1744
+ }):
1745
+ with Checkpoint() as cp:
1746
+ latest = cp.list(
1747
+ task=run["start"].task
1748
+ )[0]
1749
+ print(latest)
1750
+ cp.load(
1751
+ latest,
1752
+ "test-checkpoints"
1753
+ )
1754
+
1755
+ task = Task("TorchTuneFlow/8484/train/53673")
1756
+ with artifact_store_from(run=run, config={
1757
+ "client_params": {
1758
+ "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1759
+ "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1760
+ },
1761
+ }):
1762
+ load_model(
1763
+ task.data.model_ref,
1764
+ "test-models"
1765
+ )
1766
+ ```
1767
+ Parameters:
1533
1768
  ----------
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.
1538
- options : Dict[str, Any], default {}
1539
- Backend-specific configuration for tuning eventing behavior.
1769
+
1770
+ type: str
1771
+ The type of the datastore. Can be one of 's3', 'gcs', 'azure' or any other supported metaflow Datastore.
1772
+
1773
+ config: dict or Callable
1774
+ Dictionary of configuration options for the datastore. The following keys are required:
1775
+ - root: The root path in the datastore where the data will be saved. (needs to be in the format expected by the datastore)
1776
+ - example: 's3://bucket-name/path/to/root'
1777
+ - example: 'gs://bucket-name/path/to/root'
1778
+ - example: 'https://myblockacc.blob.core.windows.net/metaflow/'
1779
+ - role_arn (optional): AWS IAM role to access s3 bucket (only when `type` is 's3')
1780
+ - session_vars (optional): AWS session variables to access s3 bucket (only when `type` is 's3')
1781
+ - client_params (optional): AWS client parameters to access s3 bucket (only when `type` is 's3')
1540
1782
  """
1541
1783
  ...
1542
1784
 
1543
- @typing.overload
1544
- def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1545
- ...
1546
-
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] = {}):
1785
+ def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1548
1786
  """
1549
- Specifies the event(s) that this flow depends on.
1550
-
1551
- ```
1552
- @trigger(event='foo')
1553
- ```
1554
- or
1555
- ```
1556
- @trigger(events=['foo', 'bar'])
1557
- ```
1558
-
1559
- Additionally, you can specify the parameter mappings
1560
- to map event payload to Metaflow parameters for the flow.
1561
- ```
1562
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1563
- ```
1564
- or
1565
- ```
1566
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1567
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1568
- ```
1787
+ Specifies what flows belong to the same project.
1569
1788
 
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:
1575
- ```
1576
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1577
- ```
1789
+ A project-specific namespace is created for all flows that
1790
+ use the same `@project(name)`.
1578
1791
 
1579
1792
 
1580
1793
  Parameters
1581
1794
  ----------
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.
1586
- options : Dict[str, Any], default {}
1587
- Backend-specific configuration for tuning eventing behavior.
1795
+ name : str
1796
+ Project name. Make sure that the name is unique amongst all
1797
+ projects that use the same production scheduler. The name may
1798
+ contain only lowercase alphanumeric characters and underscores.
1799
+
1800
+ branch : Optional[str], default None
1801
+ The branch to use. If not specified, the branch is set to
1802
+ `user.<username>` unless `production` is set to `True`. This can
1803
+ also be set on the command line using `--branch` as a top-level option.
1804
+ It is an error to specify `branch` in the decorator and on the command line.
1805
+
1806
+ production : bool, default False
1807
+ Whether or not the branch is the production branch. This can also be set on the
1808
+ command line using `--production` as a top-level option. It is an error to specify
1809
+ `production` in the decorator and on the command line.
1810
+ The project branch name will be:
1811
+ - if `branch` is specified:
1812
+ - if `production` is True: `prod.<branch>`
1813
+ - if `production` is False: `test.<branch>`
1814
+ - if `branch` is not specified:
1815
+ - if `production` is True: `prod`
1816
+ - if `production` is False: `user.<username>`
1588
1817
  """
1589
1818
  ...
1590
1819