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