ob-metaflow-stubs 6.0.3.174__py2.py3-none-any.whl → 6.0.3.175__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 (213) hide show
  1. metaflow-stubs/__init__.pyi +643 -637
  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 +3 -3
  12. metaflow-stubs/generated_for.txt +1 -1
  13. metaflow-stubs/includefile.pyi +3 -3
  14. metaflow-stubs/info_file.pyi +1 -1
  15. metaflow-stubs/metadata_provider/__init__.pyi +1 -1
  16. metaflow-stubs/metadata_provider/heartbeat.pyi +1 -1
  17. metaflow-stubs/metadata_provider/metadata.pyi +1 -1
  18. metaflow-stubs/metadata_provider/util.pyi +1 -1
  19. metaflow-stubs/metaflow_config.pyi +1 -1
  20. metaflow-stubs/metaflow_current.pyi +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 +2 -2
  28. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/deco_injection_mixin.pyi +1 -1
  29. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/extra_components.pyi +2 -2
  30. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/__init__.pyi +1 -1
  31. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/__init__.pyi +1 -1
  32. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/checkpoint_lister.pyi +4 -4
  33. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/lineage_card.pyi +1 -1
  34. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/checkpoint_storage.pyi +3 -3
  35. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/constructors.pyi +1 -1
  36. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/core.pyi +3 -3
  37. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/decorator.pyi +3 -3
  38. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/exceptions.pyi +1 -1
  39. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/final_api.pyi +2 -2
  40. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/lineage.pyi +1 -1
  41. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/__init__.pyi +1 -1
  42. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/context.pyi +2 -2
  43. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/core.pyi +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 +2 -2
  47. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/utils.pyi +1 -1
  48. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastructures.pyi +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 +2 -2
  52. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/__init__.pyi +1 -1
  53. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/core.pyi +2 -2
  54. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/exceptions.pyi +1 -1
  55. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/model_storage.pyi +3 -3
  56. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/__init__.pyi +1 -1
  57. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/flowspec_utils.pyi +1 -1
  58. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/general.pyi +1 -1
  59. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/identity_utils.pyi +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 +1 -1
  63. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/tar_utils.pyi +1 -1
  64. metaflow-stubs/mf_extensions/outerbounds/__init__.pyi +1 -1
  65. metaflow-stubs/mf_extensions/outerbounds/plugins/__init__.pyi +1 -1
  66. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/__init__.pyi +1 -1
  67. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/async_cards.pyi +60 -0
  68. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/injector.pyi +1 -1
  69. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/__init__.pyi +1 -1
  70. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/coreweave.pyi +1 -1
  71. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/nebius.pyi +1 -1
  72. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/__init__.pyi +1 -1
  73. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/pod_killer.pyi +1 -1
  74. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/__init__.pyi +1 -1
  75. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/constants.pyi +1 -1
  76. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/exceptions.pyi +1 -1
  77. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/ollama.pyi +60 -2
  78. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/status_card.pyi +73 -0
  79. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/__init__.pyi +1 -1
  80. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/snowflake.pyi +1 -1
  81. metaflow-stubs/mf_extensions/outerbounds/profilers/__init__.pyi +1 -1
  82. metaflow-stubs/mf_extensions/outerbounds/profilers/gpu.pyi +1 -1
  83. metaflow-stubs/mf_extensions/outerbounds/remote_config.pyi +1 -1
  84. metaflow-stubs/mf_extensions/outerbounds/toplevel/__init__.pyi +1 -1
  85. metaflow-stubs/mf_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.pyi +1 -1
  86. metaflow-stubs/multicore_utils.pyi +1 -1
  87. metaflow-stubs/parameters.pyi +3 -3
  88. metaflow-stubs/plugins/__init__.pyi +9 -9
  89. metaflow-stubs/plugins/airflow/__init__.pyi +1 -1
  90. metaflow-stubs/plugins/airflow/airflow_utils.pyi +1 -1
  91. metaflow-stubs/plugins/airflow/exception.pyi +1 -1
  92. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +1 -1
  93. metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi +1 -1
  94. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +1 -1
  95. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +1 -1
  96. metaflow-stubs/plugins/argo/__init__.pyi +1 -1
  97. metaflow-stubs/plugins/argo/argo_client.pyi +1 -1
  98. metaflow-stubs/plugins/argo/argo_events.pyi +1 -1
  99. metaflow-stubs/plugins/argo/argo_workflows.pyi +2 -2
  100. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +3 -3
  101. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +1 -1
  102. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +2 -2
  103. metaflow-stubs/plugins/aws/__init__.pyi +1 -1
  104. metaflow-stubs/plugins/aws/aws_client.pyi +1 -1
  105. metaflow-stubs/plugins/aws/aws_utils.pyi +1 -1
  106. metaflow-stubs/plugins/aws/batch/__init__.pyi +1 -1
  107. metaflow-stubs/plugins/aws/batch/batch.pyi +1 -1
  108. metaflow-stubs/plugins/aws/batch/batch_client.pyi +1 -1
  109. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +1 -1
  110. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +1 -1
  111. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +4 -4
  112. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +1 -1
  113. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +1 -1
  114. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +1 -1
  115. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +1 -1
  116. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +1 -1
  117. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer.pyi +1 -1
  118. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer_objects.pyi +1 -1
  119. metaflow-stubs/plugins/azure/__init__.pyi +1 -1
  120. metaflow-stubs/plugins/azure/azure_credential.pyi +1 -1
  121. metaflow-stubs/plugins/azure/azure_exceptions.pyi +1 -1
  122. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +4 -4
  123. metaflow-stubs/plugins/azure/azure_utils.pyi +1 -1
  124. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +1 -1
  125. metaflow-stubs/plugins/azure/includefile_support.pyi +1 -1
  126. metaflow-stubs/plugins/cards/__init__.pyi +1 -1
  127. metaflow-stubs/plugins/cards/card_client.pyi +1 -1
  128. metaflow-stubs/plugins/cards/card_creator.pyi +1 -1
  129. metaflow-stubs/plugins/cards/card_datastore.pyi +1 -1
  130. metaflow-stubs/plugins/cards/card_decorator.pyi +1 -1
  131. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +1 -1
  132. metaflow-stubs/plugins/cards/card_modules/basic.pyi +2 -2
  133. metaflow-stubs/plugins/cards/card_modules/card.pyi +1 -1
  134. metaflow-stubs/plugins/cards/card_modules/components.pyi +2 -2
  135. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +1 -1
  136. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +1 -1
  137. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +1 -1
  138. metaflow-stubs/plugins/cards/card_resolver.pyi +1 -1
  139. metaflow-stubs/plugins/cards/component_serializer.pyi +1 -1
  140. metaflow-stubs/plugins/cards/exception.pyi +1 -1
  141. metaflow-stubs/plugins/catch_decorator.pyi +2 -2
  142. metaflow-stubs/plugins/datatools/__init__.pyi +1 -1
  143. metaflow-stubs/plugins/datatools/local.pyi +1 -1
  144. metaflow-stubs/plugins/datatools/s3/__init__.pyi +1 -1
  145. metaflow-stubs/plugins/datatools/s3/s3.pyi +3 -3
  146. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +1 -1
  147. metaflow-stubs/plugins/datatools/s3/s3util.pyi +1 -1
  148. metaflow-stubs/plugins/debug_logger.pyi +1 -1
  149. metaflow-stubs/plugins/debug_monitor.pyi +1 -1
  150. metaflow-stubs/plugins/environment_decorator.pyi +1 -1
  151. metaflow-stubs/plugins/events_decorator.pyi +1 -1
  152. metaflow-stubs/plugins/frameworks/__init__.pyi +1 -1
  153. metaflow-stubs/plugins/frameworks/pytorch.pyi +1 -1
  154. metaflow-stubs/plugins/gcp/__init__.pyi +1 -1
  155. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +4 -4
  156. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +1 -1
  157. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +1 -1
  158. metaflow-stubs/plugins/gcp/gs_utils.pyi +1 -1
  159. metaflow-stubs/plugins/gcp/includefile_support.pyi +1 -1
  160. metaflow-stubs/plugins/kubernetes/__init__.pyi +1 -1
  161. metaflow-stubs/plugins/kubernetes/kube_utils.pyi +1 -1
  162. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +1 -1
  163. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +1 -1
  164. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +1 -1
  165. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +1 -1
  166. metaflow-stubs/plugins/kubernetes/spot_monitor_sidecar.pyi +1 -1
  167. metaflow-stubs/plugins/ollama/__init__.pyi +13 -4
  168. metaflow-stubs/plugins/parallel_decorator.pyi +1 -1
  169. metaflow-stubs/plugins/perimeters.pyi +1 -1
  170. metaflow-stubs/plugins/project_decorator.pyi +1 -1
  171. metaflow-stubs/plugins/pypi/__init__.pyi +2 -2
  172. metaflow-stubs/plugins/pypi/conda_decorator.pyi +1 -1
  173. metaflow-stubs/plugins/pypi/conda_environment.pyi +5 -5
  174. metaflow-stubs/plugins/pypi/parsers.pyi +1 -1
  175. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +1 -1
  176. metaflow-stubs/plugins/pypi/pypi_environment.pyi +1 -1
  177. metaflow-stubs/plugins/pypi/utils.pyi +1 -1
  178. metaflow-stubs/plugins/resources_decorator.pyi +1 -1
  179. metaflow-stubs/plugins/retry_decorator.pyi +1 -1
  180. metaflow-stubs/plugins/secrets/__init__.pyi +2 -2
  181. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +3 -3
  182. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +1 -1
  183. metaflow-stubs/plugins/snowflake/__init__.pyi +1 -1
  184. metaflow-stubs/plugins/storage_executor.pyi +1 -1
  185. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +2 -2
  186. metaflow-stubs/plugins/timeout_decorator.pyi +2 -2
  187. metaflow-stubs/plugins/torchtune/__init__.pyi +1 -1
  188. metaflow-stubs/plugins/uv/__init__.pyi +1 -1
  189. metaflow-stubs/plugins/uv/uv_environment.pyi +2 -2
  190. metaflow-stubs/profilers/__init__.pyi +1 -1
  191. metaflow-stubs/pylint_wrapper.pyi +1 -1
  192. metaflow-stubs/runner/__init__.pyi +1 -1
  193. metaflow-stubs/runner/deployer.pyi +4 -4
  194. metaflow-stubs/runner/deployer_impl.pyi +2 -2
  195. metaflow-stubs/runner/metaflow_runner.pyi +2 -2
  196. metaflow-stubs/runner/nbdeploy.pyi +1 -1
  197. metaflow-stubs/runner/nbrun.pyi +1 -1
  198. metaflow-stubs/runner/subprocess_manager.pyi +1 -1
  199. metaflow-stubs/runner/utils.pyi +2 -2
  200. metaflow-stubs/system/__init__.pyi +1 -1
  201. metaflow-stubs/system/system_logger.pyi +1 -1
  202. metaflow-stubs/system/system_monitor.pyi +1 -1
  203. metaflow-stubs/tagging_util.pyi +1 -1
  204. metaflow-stubs/tuple_util.pyi +1 -1
  205. metaflow-stubs/user_configs/__init__.pyi +1 -1
  206. metaflow-stubs/user_configs/config_decorators.pyi +5 -5
  207. metaflow-stubs/user_configs/config_options.pyi +2 -2
  208. metaflow-stubs/user_configs/config_parameters.pyi +6 -6
  209. {ob_metaflow_stubs-6.0.3.174.dist-info → ob_metaflow_stubs-6.0.3.175.dist-info}/METADATA +1 -1
  210. ob_metaflow_stubs-6.0.3.175.dist-info/RECORD +213 -0
  211. ob_metaflow_stubs-6.0.3.174.dist-info/RECORD +0 -211
  212. {ob_metaflow_stubs-6.0.3.174.dist-info → ob_metaflow_stubs-6.0.3.175.dist-info}/WHEEL +0 -0
  213. {ob_metaflow_stubs-6.0.3.174.dist-info → ob_metaflow_stubs-6.0.3.175.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-29T18:56:59.040642 #
4
+ # Generated on 2025-06-02T15:33:51.427981 #
5
5
  ######################################################################################################
6
6
 
7
7
  from __future__ import annotations
@@ -44,8 +44,8 @@ from . import plugins as plugins
44
44
  from .mf_extensions.outerbounds.toplevel.global_aliases_for_metaflow_package import S3 as S3
45
45
  from . import includefile as includefile
46
46
  from .includefile import IncludeFile as IncludeFile
47
- from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
48
47
  from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
48
+ from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
49
49
  from .plugins.pypi.parsers import conda_environment_yml_parser as conda_environment_yml_parser
50
50
  from . import client as client
51
51
  from .client.core import namespace as namespace
@@ -154,479 +154,507 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
154
154
  ...
155
155
 
156
156
  @typing.overload
157
- def resources(*, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
157
+ def checkpoint(*, load_policy: str = 'fresh', temp_dir_root: str = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
158
158
  """
159
- Specifies the resources needed when executing this step.
160
-
161
- Use `@resources` to specify the resource requirements
162
- independently of the specific compute layer (`@batch`, `@kubernetes`).
159
+ Enables checkpointing for a step.
163
160
 
164
- You can choose the compute layer on the command line by executing e.g.
165
- ```
166
- python myflow.py run --with batch
167
- ```
168
- or
169
- ```
170
- python myflow.py run --with kubernetes
171
- ```
172
- which executes the flow on the desired system using the
173
- requirements specified in `@resources`.
174
161
 
175
162
 
176
163
  Parameters
177
164
  ----------
178
- cpu : int, default 1
179
- Number of CPUs required for this step.
180
- gpu : int, optional, default None
181
- Number of GPUs required for this step.
182
- disk : int, optional, default None
183
- Disk size (in MB) required for this step. Only applies on Kubernetes.
184
- memory : int, default 4096
185
- Memory size (in MB) required for this step.
186
- shared_memory : int, optional, default None
187
- The value for the size (in MiB) of the /dev/shm volume for this step.
188
- This parameter maps to the `--shm-size` option in Docker.
165
+ load_policy : str, default: "fresh"
166
+ The policy for loading the checkpoint. The following policies are supported:
167
+ - "eager": Loads the the latest available checkpoint within the namespace.
168
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
169
+ will be loaded at the start of the task.
170
+ - "none": Do not load any checkpoint
171
+ - "fresh": Loads the lastest checkpoint created within the running Task.
172
+ This mode helps loading checkpoints across various retry attempts of the same task.
173
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
174
+ created within the task will be loaded when the task is retries execution on failure.
175
+
176
+ temp_dir_root : str, default: None
177
+ The root directory under which `current.checkpoint.directory` will be created.
189
178
  """
190
179
  ...
191
180
 
192
181
  @typing.overload
193
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
182
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
194
183
  ...
195
184
 
196
185
  @typing.overload
197
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
186
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
198
187
  ...
199
188
 
200
- def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None):
189
+ def checkpoint(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, load_policy: str = 'fresh', temp_dir_root: str = None):
201
190
  """
202
- Specifies the resources needed when executing this step.
203
-
204
- Use `@resources` to specify the resource requirements
205
- independently of the specific compute layer (`@batch`, `@kubernetes`).
191
+ Enables checkpointing for a step.
206
192
 
207
- You can choose the compute layer on the command line by executing e.g.
208
- ```
209
- python myflow.py run --with batch
210
- ```
211
- or
212
- ```
213
- python myflow.py run --with kubernetes
214
- ```
215
- which executes the flow on the desired system using the
216
- requirements specified in `@resources`.
217
193
 
218
194
 
219
195
  Parameters
220
196
  ----------
221
- cpu : int, default 1
222
- Number of CPUs required for this step.
223
- gpu : int, optional, default None
224
- Number of GPUs required for this step.
225
- disk : int, optional, default None
226
- Disk size (in MB) required for this step. Only applies on Kubernetes.
227
- memory : int, default 4096
228
- Memory size (in MB) required for this step.
229
- shared_memory : int, optional, default None
230
- The value for the size (in MiB) of the /dev/shm volume for this step.
231
- This parameter maps to the `--shm-size` option in Docker.
197
+ load_policy : str, default: "fresh"
198
+ The policy for loading the checkpoint. The following policies are supported:
199
+ - "eager": Loads the the latest available checkpoint within the namespace.
200
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
201
+ will be loaded at the start of the task.
202
+ - "none": Do not load any checkpoint
203
+ - "fresh": Loads the lastest checkpoint created within the running Task.
204
+ This mode helps loading checkpoints across various retry attempts of the same task.
205
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
206
+ created within the task will be loaded when the task is retries execution on failure.
207
+
208
+ temp_dir_root : str, default: None
209
+ The root directory under which `current.checkpoint.directory` will be created.
232
210
  """
233
211
  ...
234
212
 
235
213
  @typing.overload
236
- 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]]]:
214
+ 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]]]:
237
215
  """
238
- Specifies a timeout for your step.
239
-
240
- This decorator is useful if this step may hang indefinitely.
241
-
242
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
243
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
244
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
216
+ Specifies the Conda environment for the step.
245
217
 
246
- Note that all the values specified in parameters are added together so if you specify
247
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
218
+ Information in this decorator will augment any
219
+ attributes set in the `@conda_base` flow-level decorator. Hence,
220
+ you can use `@conda_base` to set packages required by all
221
+ steps and use `@conda` to specify step-specific overrides.
248
222
 
249
223
 
250
224
  Parameters
251
225
  ----------
252
- seconds : int, default 0
253
- Number of seconds to wait prior to timing out.
254
- minutes : int, default 0
255
- Number of minutes to wait prior to timing out.
256
- hours : int, default 0
257
- Number of hours to wait prior to timing out.
226
+ packages : Dict[str, str], default {}
227
+ Packages to use for this step. The key is the name of the package
228
+ and the value is the version to use.
229
+ libraries : Dict[str, str], default {}
230
+ Supported for backward compatibility. When used with packages, packages will take precedence.
231
+ python : str, optional, default None
232
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
233
+ that the version used will correspond to the version of the Python interpreter used to start the run.
234
+ disabled : bool, default False
235
+ If set to True, disables @conda.
258
236
  """
259
237
  ...
260
238
 
261
239
  @typing.overload
262
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
240
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
263
241
  ...
264
242
 
265
243
  @typing.overload
266
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
244
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
267
245
  ...
268
246
 
269
- 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):
247
+ 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):
270
248
  """
271
- Specifies a timeout for your step.
272
-
273
- This decorator is useful if this step may hang indefinitely.
274
-
275
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
276
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
277
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
249
+ Specifies the Conda environment for the step.
278
250
 
279
- Note that all the values specified in parameters are added together so if you specify
280
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
251
+ Information in this decorator will augment any
252
+ attributes set in the `@conda_base` flow-level decorator. Hence,
253
+ you can use `@conda_base` to set packages required by all
254
+ steps and use `@conda` to specify step-specific overrides.
281
255
 
282
256
 
283
257
  Parameters
284
258
  ----------
285
- seconds : int, default 0
286
- Number of seconds to wait prior to timing out.
287
- minutes : int, default 0
288
- Number of minutes to wait prior to timing out.
289
- hours : int, default 0
290
- Number of hours to wait prior to timing out.
259
+ packages : Dict[str, str], default {}
260
+ Packages to use for this step. The key is the name of the package
261
+ and the value is the version to use.
262
+ libraries : Dict[str, str], default {}
263
+ Supported for backward compatibility. When used with packages, packages will take precedence.
264
+ python : str, optional, default None
265
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
266
+ that the version used will correspond to the version of the Python interpreter used to start the run.
267
+ disabled : bool, default False
268
+ If set to True, disables @conda.
291
269
  """
292
270
  ...
293
271
 
294
- 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]]]:
272
+ @typing.overload
273
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
295
274
  """
296
- This decorator is used to run Ollama APIs as Metaflow task sidecars.
297
-
298
- User code call
299
- --------------
300
- @ollama(
301
- models=[...],
302
- ...
303
- )
304
-
305
- Valid backend options
306
- ---------------------
307
- - 'local': Run as a separate process on the local task machine.
308
- - (TODO) 'managed': Outerbounds hosts and selects compute provider.
309
- - (TODO) 'remote': Spin up separate instance to serve Ollama models.
310
-
311
- Valid model options
312
- -------------------
313
- Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
314
-
315
-
316
- Parameters
317
- ----------
318
- models: list[str]
319
- List of Ollama containers running models in sidecars.
320
- backend: str
321
- Determines where and how to run the Ollama process.
322
- force_pull: bool
323
- Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
324
- skip_push_check: bool
325
- Whether to skip the check that populates/overwrites remote cache on terminating an ollama model.
326
- debug: bool
327
- Whether to turn on verbose debugging logs.
275
+ Internal decorator to support Fast bakery
328
276
  """
329
277
  ...
330
278
 
331
- 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]]]:
279
+ @typing.overload
280
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
281
+ ...
282
+
283
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
332
284
  """
333
- Decorator that helps cache, version and store models/datasets from huggingface hub.
334
-
335
-
336
- Parameters
337
- ----------
338
- temp_dir_root : str, optional
339
- The root directory that will hold the temporary directory where objects will be downloaded.
340
-
341
- load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
342
- The list of repos (models/datasets) to load.
343
-
344
- Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
345
-
346
- - If repo (model/dataset) is not found in the datastore:
347
- - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
348
- - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
349
- - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
350
-
351
- - If repo is found in the datastore:
352
- - Loads it directly from datastore to local path (can be temporary directory or specified path)
285
+ Internal decorator to support Fast bakery
353
286
  """
354
287
  ...
355
288
 
356
289
  @typing.overload
357
- 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]]]:
290
+ 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]]]:
358
291
  """
359
- Specifies the number of times the task corresponding
360
- to a step needs to be retried.
361
-
362
- This decorator is useful for handling transient errors, such as networking issues.
363
- If your task contains operations that can't be retried safely, e.g. database updates,
364
- it is advisable to annotate it with `@retry(times=0)`.
292
+ Specifies that the step will success under all circumstances.
365
293
 
366
- This can be used in conjunction with the `@catch` decorator. The `@catch`
367
- decorator will execute a no-op task after all retries have been exhausted,
368
- ensuring that the flow execution can continue.
294
+ The decorator will create an optional artifact, specified by `var`, which
295
+ contains the exception raised. You can use it to detect the presence
296
+ of errors, indicating that all happy-path artifacts produced by the step
297
+ are missing.
369
298
 
370
299
 
371
300
  Parameters
372
301
  ----------
373
- times : int, default 3
374
- Number of times to retry this task.
375
- minutes_between_retries : int, default 2
376
- Number of minutes between retries.
302
+ var : str, optional, default None
303
+ Name of the artifact in which to store the caught exception.
304
+ If not specified, the exception is not stored.
305
+ print_exception : bool, default True
306
+ Determines whether or not the exception is printed to
307
+ stdout when caught.
377
308
  """
378
309
  ...
379
310
 
380
311
  @typing.overload
381
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
312
+ def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
382
313
  ...
383
314
 
384
315
  @typing.overload
385
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
316
+ def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
386
317
  ...
387
318
 
388
- 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):
319
+ 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):
389
320
  """
390
- Specifies the number of times the task corresponding
391
- to a step needs to be retried.
392
-
393
- This decorator is useful for handling transient errors, such as networking issues.
394
- If your task contains operations that can't be retried safely, e.g. database updates,
395
- it is advisable to annotate it with `@retry(times=0)`.
321
+ Specifies that the step will success under all circumstances.
396
322
 
397
- This can be used in conjunction with the `@catch` decorator. The `@catch`
398
- decorator will execute a no-op task after all retries have been exhausted,
399
- ensuring that the flow execution can continue.
323
+ The decorator will create an optional artifact, specified by `var`, which
324
+ contains the exception raised. You can use it to detect the presence
325
+ of errors, indicating that all happy-path artifacts produced by the step
326
+ are missing.
400
327
 
401
328
 
402
329
  Parameters
403
330
  ----------
404
- times : int, default 3
405
- Number of times to retry this task.
406
- minutes_between_retries : int, default 2
407
- Number of minutes between retries.
331
+ var : str, optional, default None
332
+ Name of the artifact in which to store the caught exception.
333
+ If not specified, the exception is not stored.
334
+ print_exception : bool, default True
335
+ Determines whether or not the exception is printed to
336
+ stdout when caught.
408
337
  """
409
338
  ...
410
339
 
411
340
  @typing.overload
412
- 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]]]:
341
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
413
342
  """
414
- Specifies environment variables to be set prior to the execution of a step.
343
+ Decorator prototype for all step decorators. This function gets specialized
344
+ and imported for all decorators types by _import_plugin_decorators().
345
+ """
346
+ ...
347
+
348
+ @typing.overload
349
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
350
+ ...
351
+
352
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
353
+ """
354
+ Decorator prototype for all step decorators. This function gets specialized
355
+ and imported for all decorators types by _import_plugin_decorators().
356
+ """
357
+ ...
358
+
359
+ @typing.overload
360
+ 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]]]:
361
+ """
362
+ Specifies the number of times the task corresponding
363
+ to a step needs to be retried.
364
+
365
+ This decorator is useful for handling transient errors, such as networking issues.
366
+ If your task contains operations that can't be retried safely, e.g. database updates,
367
+ it is advisable to annotate it with `@retry(times=0)`.
368
+
369
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
370
+ decorator will execute a no-op task after all retries have been exhausted,
371
+ ensuring that the flow execution can continue.
415
372
 
416
373
 
417
374
  Parameters
418
375
  ----------
419
- vars : Dict[str, str], default {}
420
- Dictionary of environment variables to set.
376
+ times : int, default 3
377
+ Number of times to retry this task.
378
+ minutes_between_retries : int, default 2
379
+ Number of minutes between retries.
421
380
  """
422
381
  ...
423
382
 
424
383
  @typing.overload
425
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
384
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
426
385
  ...
427
386
 
428
387
  @typing.overload
429
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
388
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
430
389
  ...
431
390
 
432
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
391
+ 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):
433
392
  """
434
- Specifies environment variables to be set prior to the execution of a step.
393
+ Specifies the number of times the task corresponding
394
+ to a step needs to be retried.
395
+
396
+ This decorator is useful for handling transient errors, such as networking issues.
397
+ If your task contains operations that can't be retried safely, e.g. database updates,
398
+ it is advisable to annotate it with `@retry(times=0)`.
399
+
400
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
401
+ decorator will execute a no-op task after all retries have been exhausted,
402
+ ensuring that the flow execution can continue.
435
403
 
436
404
 
437
405
  Parameters
438
406
  ----------
439
- vars : Dict[str, str], default {}
440
- Dictionary of environment variables to set.
407
+ times : int, default 3
408
+ Number of times to retry this task.
409
+ minutes_between_retries : int, default 2
410
+ Number of minutes between retries.
441
411
  """
442
412
  ...
443
413
 
444
- 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]]]:
414
+ def ollama(*, models: list, backend: str, force_pull: bool, cache_update_policy: str, force_cache_update: bool, debug: bool, circuit_breaker_config: dict, timeout_config: dict) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
445
415
  """
446
- Specifies that this step is used to deploy an instance of the app.
447
- Requires that self.app_name, self.app_port, self.entrypoint and self.deployDir is set.
416
+ This decorator is used to run Ollama APIs as Metaflow task sidecars.
417
+
418
+ User code call
419
+ --------------
420
+ @ollama(
421
+ models=[...],
422
+ ...
423
+ )
424
+
425
+ Valid backend options
426
+ ---------------------
427
+ - 'local': Run as a separate process on the local task machine.
428
+ - (TODO) 'managed': Outerbounds hosts and selects compute provider.
429
+ - (TODO) 'remote': Spin up separate instance to serve Ollama models.
430
+
431
+ Valid model options
432
+ -------------------
433
+ Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
448
434
 
449
435
 
450
436
  Parameters
451
437
  ----------
452
- app_port : int
453
- Number of GPUs to use.
454
- app_name : str
455
- Name of the app to deploy.
438
+ models: list[str]
439
+ List of Ollama containers running models in sidecars.
440
+ backend: str
441
+ Determines where and how to run the Ollama process.
442
+ force_pull: bool
443
+ Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
444
+ cache_update_policy: str
445
+ Cache update policy: "auto", "force", or "never".
446
+ force_cache_update: bool
447
+ Simple override for "force" cache update policy.
448
+ debug: bool
449
+ Whether to turn on verbose debugging logs.
450
+ circuit_breaker_config: dict
451
+ Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
452
+ timeout_config: dict
453
+ Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
456
454
  """
457
455
  ...
458
456
 
459
- 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]]]:
457
+ 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]]]:
460
458
  """
461
- Specifies that this step should execute on DGX cloud.
459
+ Decorator that helps cache, version and store models/datasets from huggingface hub.
462
460
 
463
461
 
464
462
  Parameters
465
463
  ----------
466
- gpu : int
467
- Number of GPUs to use.
468
- gpu_type : str
469
- Type of Nvidia GPU to use.
470
- """
471
- ...
472
-
473
- @typing.overload
474
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
475
- """
476
- Decorator prototype for all step decorators. This function gets specialized
477
- and imported for all decorators types by _import_plugin_decorators().
478
- """
479
- ...
480
-
481
- @typing.overload
482
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
483
- ...
484
-
485
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
486
- """
487
- Decorator prototype for all step decorators. This function gets specialized
488
- and imported for all decorators types by _import_plugin_decorators().
464
+ temp_dir_root : str, optional
465
+ The root directory that will hold the temporary directory where objects will be downloaded.
466
+
467
+ load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
468
+ The list of repos (models/datasets) to load.
469
+
470
+ Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
471
+
472
+ - If repo (model/dataset) is not found in the datastore:
473
+ - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
474
+ - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
475
+ - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
476
+
477
+ - If repo is found in the datastore:
478
+ - Loads it directly from datastore to local path (can be temporary directory or specified path)
489
479
  """
490
480
  ...
491
481
 
492
482
  @typing.overload
493
- 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]]]:
483
+ 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]]]:
494
484
  """
495
- Creates a human-readable report, a Metaflow Card, after this step completes.
485
+ Specifies a timeout for your step.
496
486
 
497
- Note that you may add multiple `@card` decorators in a step with different parameters.
487
+ This decorator is useful if this step may hang indefinitely.
488
+
489
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
490
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
491
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
492
+
493
+ Note that all the values specified in parameters are added together so if you specify
494
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
498
495
 
499
496
 
500
497
  Parameters
501
498
  ----------
502
- type : str, default 'default'
503
- Card type.
504
- id : str, optional, default None
505
- If multiple cards are present, use this id to identify this card.
506
- options : Dict[str, Any], default {}
507
- Options passed to the card. The contents depend on the card type.
508
- timeout : int, default 45
509
- Interrupt reporting if it takes more than this many seconds.
499
+ seconds : int, default 0
500
+ Number of seconds to wait prior to timing out.
501
+ minutes : int, default 0
502
+ Number of minutes to wait prior to timing out.
503
+ hours : int, default 0
504
+ Number of hours to wait prior to timing out.
510
505
  """
511
506
  ...
512
507
 
513
508
  @typing.overload
514
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
509
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
515
510
  ...
516
511
 
517
512
  @typing.overload
518
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
513
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
519
514
  ...
520
515
 
521
- 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):
516
+ 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):
522
517
  """
523
- Creates a human-readable report, a Metaflow Card, after this step completes.
518
+ Specifies a timeout for your step.
524
519
 
525
- Note that you may add multiple `@card` decorators in a step with different parameters.
520
+ This decorator is useful if this step may hang indefinitely.
521
+
522
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
523
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
524
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
525
+
526
+ Note that all the values specified in parameters are added together so if you specify
527
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
526
528
 
527
529
 
528
530
  Parameters
529
531
  ----------
530
- type : str, default 'default'
531
- Card type.
532
- id : str, optional, default None
533
- If multiple cards are present, use this id to identify this card.
534
- options : Dict[str, Any], default {}
535
- Options passed to the card. The contents depend on the card type.
536
- timeout : int, default 45
537
- Interrupt reporting if it takes more than this many seconds.
532
+ seconds : int, default 0
533
+ Number of seconds to wait prior to timing out.
534
+ minutes : int, default 0
535
+ Number of minutes to wait prior to timing out.
536
+ hours : int, default 0
537
+ Number of hours to wait prior to timing out.
538
538
  """
539
539
  ...
540
540
 
541
541
  @typing.overload
542
- def checkpoint(*, load_policy: str = 'fresh', temp_dir_root: str = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
542
+ 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]]]:
543
543
  """
544
- Enables checkpointing for a step.
544
+ Enables loading / saving of models within a step.
545
545
 
546
546
 
547
547
 
548
548
  Parameters
549
549
  ----------
550
- load_policy : str, default: "fresh"
551
- The policy for loading the checkpoint. The following policies are supported:
552
- - "eager": Loads the the latest available checkpoint within the namespace.
553
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
554
- will be loaded at the start of the task.
555
- - "none": Do not load any checkpoint
556
- - "fresh": Loads the lastest checkpoint created within the running Task.
557
- This mode helps loading checkpoints across various retry attempts of the same task.
558
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
559
- created within the task will be loaded when the task is retries execution on failure.
550
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
551
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
552
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
553
+ - `current.checkpoint`
554
+ - `current.model`
555
+ - `current.huggingface_hub`
556
+
557
+ 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
558
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
559
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
560
560
 
561
561
  temp_dir_root : str, default: None
562
- The root directory under which `current.checkpoint.directory` will be created.
562
+ The root directory under which `current.model.loaded` will store loaded models
563
563
  """
564
564
  ...
565
565
 
566
566
  @typing.overload
567
- def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
567
+ def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
568
568
  ...
569
569
 
570
570
  @typing.overload
571
- def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
571
+ def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
572
572
  ...
573
573
 
574
- def checkpoint(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, load_policy: str = 'fresh', temp_dir_root: str = None):
574
+ 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):
575
575
  """
576
- Enables checkpointing for a step.
576
+ Enables loading / saving of models within a step.
577
577
 
578
578
 
579
579
 
580
580
  Parameters
581
581
  ----------
582
- load_policy : str, default: "fresh"
583
- The policy for loading the checkpoint. The following policies are supported:
584
- - "eager": Loads the the latest available checkpoint within the namespace.
585
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
586
- will be loaded at the start of the task.
587
- - "none": Do not load any checkpoint
588
- - "fresh": Loads the lastest checkpoint created within the running Task.
589
- This mode helps loading checkpoints across various retry attempts of the same task.
590
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
591
- created within the task will be loaded when the task is retries execution on failure.
582
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
583
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
584
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
585
+ - `current.checkpoint`
586
+ - `current.model`
587
+ - `current.huggingface_hub`
588
+
589
+ 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
590
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
591
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
592
592
 
593
593
  temp_dir_root : str, default: None
594
- The root directory under which `current.checkpoint.directory` will be created.
594
+ The root directory under which `current.model.loaded` will store loaded models
595
595
  """
596
596
  ...
597
597
 
598
598
  @typing.overload
599
- 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]]]:
599
+ 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]]]:
600
600
  """
601
- Specifies secrets to be retrieved and injected as environment variables prior to
602
- the execution of a step.
601
+ Creates a human-readable report, a Metaflow Card, after this step completes.
602
+
603
+ Note that you may add multiple `@card` decorators in a step with different parameters.
603
604
 
604
605
 
605
606
  Parameters
606
607
  ----------
607
- sources : List[Union[str, Dict[str, Any]]], default: []
608
- List of secret specs, defining how the secrets are to be retrieved
608
+ type : str, default 'default'
609
+ Card type.
610
+ id : str, optional, default None
611
+ If multiple cards are present, use this id to identify this card.
612
+ options : Dict[str, Any], default {}
613
+ Options passed to the card. The contents depend on the card type.
614
+ timeout : int, default 45
615
+ Interrupt reporting if it takes more than this many seconds.
609
616
  """
610
617
  ...
611
618
 
612
619
  @typing.overload
613
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
620
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
614
621
  ...
615
622
 
616
623
  @typing.overload
617
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
624
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
618
625
  ...
619
626
 
620
- 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]]] = []):
627
+ 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):
621
628
  """
622
- Specifies secrets to be retrieved and injected as environment variables prior to
623
- the execution of a step.
629
+ Creates a human-readable report, a Metaflow Card, after this step completes.
630
+
631
+ Note that you may add multiple `@card` decorators in a step with different parameters.
624
632
 
625
633
 
626
634
  Parameters
627
635
  ----------
628
- sources : List[Union[str, Dict[str, Any]]], default: []
629
- List of secret specs, defining how the secrets are to be retrieved
636
+ type : str, default 'default'
637
+ Card type.
638
+ id : str, optional, default None
639
+ If multiple cards are present, use this id to identify this card.
640
+ options : Dict[str, Any], default {}
641
+ Options passed to the card. The contents depend on the card type.
642
+ timeout : int, default 45
643
+ Interrupt reporting if it takes more than this many seconds.
644
+ """
645
+ ...
646
+
647
+ 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]]]:
648
+ """
649
+ Specifies that this step should execute on DGX cloud.
650
+
651
+
652
+ Parameters
653
+ ----------
654
+ gpu : int
655
+ Number of GPUs to use.
656
+ gpu_type : str
657
+ Type of Nvidia GPU to use.
630
658
  """
631
659
  ...
632
660
 
@@ -681,111 +709,102 @@ def pypi(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typ
681
709
  """
682
710
  ...
683
711
 
684
- @typing.overload
685
- 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]]]:
712
+ 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]]]:
686
713
  """
687
- Specifies that the step will success under all circumstances.
714
+ Specifies that this step should execute on DGX cloud.
688
715
 
689
- The decorator will create an optional artifact, specified by `var`, which
690
- contains the exception raised. You can use it to detect the presence
691
- of errors, indicating that all happy-path artifacts produced by the step
692
- are missing.
716
+
717
+ Parameters
718
+ ----------
719
+ gpu : int
720
+ Number of GPUs to use.
721
+ gpu_type : str
722
+ Type of Nvidia GPU to use.
723
+ queue_timeout : int
724
+ Time to keep the job in NVCF's queue.
725
+ """
726
+ ...
727
+
728
+ 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]]]:
729
+ """
730
+ Specifies that this step is used to deploy an instance of the app.
731
+ Requires that self.app_name, self.app_port, self.entrypoint and self.deployDir is set.
693
732
 
694
733
 
695
734
  Parameters
696
735
  ----------
697
- var : str, optional, default None
698
- Name of the artifact in which to store the caught exception.
699
- If not specified, the exception is not stored.
700
- print_exception : bool, default True
701
- Determines whether or not the exception is printed to
702
- stdout when caught.
736
+ app_port : int
737
+ Number of GPUs to use.
738
+ app_name : str
739
+ Name of the app to deploy.
703
740
  """
704
741
  ...
705
742
 
706
743
  @typing.overload
707
- def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
744
+ 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]]]:
745
+ """
746
+ Specifies environment variables to be set prior to the execution of a step.
747
+
748
+
749
+ Parameters
750
+ ----------
751
+ vars : Dict[str, str], default {}
752
+ Dictionary of environment variables to set.
753
+ """
708
754
  ...
709
755
 
710
756
  @typing.overload
711
- def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
757
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
712
758
  ...
713
759
 
714
- 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):
760
+ @typing.overload
761
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
762
+ ...
763
+
764
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
715
765
  """
716
- Specifies that the step will success under all circumstances.
717
-
718
- The decorator will create an optional artifact, specified by `var`, which
719
- contains the exception raised. You can use it to detect the presence
720
- of errors, indicating that all happy-path artifacts produced by the step
721
- are missing.
766
+ Specifies environment variables to be set prior to the execution of a step.
722
767
 
723
768
 
724
769
  Parameters
725
770
  ----------
726
- var : str, optional, default None
727
- Name of the artifact in which to store the caught exception.
728
- If not specified, the exception is not stored.
729
- print_exception : bool, default True
730
- Determines whether or not the exception is printed to
731
- stdout when caught.
771
+ vars : Dict[str, str], default {}
772
+ Dictionary of environment variables to set.
732
773
  """
733
774
  ...
734
775
 
735
776
  @typing.overload
736
- 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]]]:
777
+ 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]]]:
737
778
  """
738
- Enables loading / saving of models within a step.
739
-
779
+ Specifies secrets to be retrieved and injected as environment variables prior to
780
+ the execution of a step.
740
781
 
741
782
 
742
783
  Parameters
743
784
  ----------
744
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
745
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
746
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
747
- - `current.checkpoint`
748
- - `current.model`
749
- - `current.huggingface_hub`
750
-
751
- 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
752
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
753
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
754
-
755
- temp_dir_root : str, default: None
756
- The root directory under which `current.model.loaded` will store loaded models
785
+ sources : List[Union[str, Dict[str, Any]]], default: []
786
+ List of secret specs, defining how the secrets are to be retrieved
757
787
  """
758
788
  ...
759
789
 
760
790
  @typing.overload
761
- def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
791
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
762
792
  ...
763
793
 
764
794
  @typing.overload
765
- def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
795
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
766
796
  ...
767
797
 
768
- 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):
798
+ 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]]] = []):
769
799
  """
770
- Enables loading / saving of models within a step.
771
-
800
+ Specifies secrets to be retrieved and injected as environment variables prior to
801
+ the execution of a step.
772
802
 
773
803
 
774
804
  Parameters
775
805
  ----------
776
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
777
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
778
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
779
- - `current.checkpoint`
780
- - `current.model`
781
- - `current.huggingface_hub`
782
-
783
- 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
784
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
785
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
786
-
787
- temp_dir_root : str, default: None
788
- The root directory under which `current.model.loaded` will store loaded models
806
+ sources : List[Union[str, Dict[str, Any]]], default: []
807
+ List of secret specs, defining how the secrets are to be retrieved
789
808
  """
790
809
  ...
791
810
 
@@ -874,196 +893,133 @@ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: ty
874
893
  """
875
894
  ...
876
895
 
877
- 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]]]:
878
- """
879
- Specifies that this step should execute on DGX cloud.
880
-
881
-
882
- Parameters
883
- ----------
884
- gpu : int
885
- Number of GPUs to use.
886
- gpu_type : str
887
- Type of Nvidia GPU to use.
888
- queue_timeout : int
889
- Time to keep the job in NVCF's queue.
890
- """
891
- ...
892
-
893
896
  @typing.overload
894
- 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]]]:
897
+ def resources(*, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
895
898
  """
896
- Specifies the Conda environment for the step.
897
-
898
- Information in this decorator will augment any
899
- attributes set in the `@conda_base` flow-level decorator. Hence,
900
- you can use `@conda_base` to set packages required by all
901
- steps and use `@conda` to specify step-specific overrides.
902
-
899
+ Specifies the resources needed when executing this step.
903
900
 
904
- Parameters
905
- ----------
906
- packages : Dict[str, str], default {}
907
- Packages to use for this step. The key is the name of the package
908
- and the value is the version to use.
909
- libraries : Dict[str, str], default {}
910
- Supported for backward compatibility. When used with packages, packages will take precedence.
911
- python : str, optional, default None
912
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
913
- that the version used will correspond to the version of the Python interpreter used to start the run.
914
- disabled : bool, default False
915
- If set to True, disables @conda.
916
- """
917
- ...
918
-
919
- @typing.overload
920
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
921
- ...
922
-
923
- @typing.overload
924
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
925
- ...
926
-
927
- 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):
928
- """
929
- Specifies the Conda environment for the step.
901
+ Use `@resources` to specify the resource requirements
902
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
930
903
 
931
- Information in this decorator will augment any
932
- attributes set in the `@conda_base` flow-level decorator. Hence,
933
- you can use `@conda_base` to set packages required by all
934
- steps and use `@conda` to specify step-specific overrides.
904
+ You can choose the compute layer on the command line by executing e.g.
905
+ ```
906
+ python myflow.py run --with batch
907
+ ```
908
+ or
909
+ ```
910
+ python myflow.py run --with kubernetes
911
+ ```
912
+ which executes the flow on the desired system using the
913
+ requirements specified in `@resources`.
935
914
 
936
915
 
937
916
  Parameters
938
917
  ----------
939
- packages : Dict[str, str], default {}
940
- Packages to use for this step. The key is the name of the package
941
- and the value is the version to use.
942
- libraries : Dict[str, str], default {}
943
- Supported for backward compatibility. When used with packages, packages will take precedence.
944
- python : str, optional, default None
945
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
946
- that the version used will correspond to the version of the Python interpreter used to start the run.
947
- disabled : bool, default False
948
- If set to True, disables @conda.
918
+ cpu : int, default 1
919
+ Number of CPUs required for this step.
920
+ gpu : int, optional, default None
921
+ Number of GPUs required for this step.
922
+ disk : int, optional, default None
923
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
924
+ memory : int, default 4096
925
+ Memory size (in MB) required for this step.
926
+ shared_memory : int, optional, default None
927
+ The value for the size (in MiB) of the /dev/shm volume for this step.
928
+ This parameter maps to the `--shm-size` option in Docker.
949
929
  """
950
930
  ...
951
931
 
952
932
  @typing.overload
953
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
954
- """
955
- Internal decorator to support Fast bakery
956
- """
933
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
957
934
  ...
958
935
 
959
936
  @typing.overload
960
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
961
- ...
962
-
963
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
964
- """
965
- Internal decorator to support Fast bakery
966
- """
937
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
967
938
  ...
968
939
 
969
- @typing.overload
970
- 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]]:
940
+ def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None):
971
941
  """
972
- Specifies the flow(s) that this flow depends on.
942
+ Specifies the resources needed when executing this step.
973
943
 
974
- ```
975
- @trigger_on_finish(flow='FooFlow')
976
- ```
977
- or
978
- ```
979
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
980
- ```
981
- This decorator respects the @project decorator and triggers the flow
982
- when upstream runs within the same namespace complete successfully
944
+ Use `@resources` to specify the resource requirements
945
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
983
946
 
984
- Additionally, you can specify project aware upstream flow dependencies
985
- by specifying the fully qualified project_flow_name.
947
+ You can choose the compute layer on the command line by executing e.g.
986
948
  ```
987
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
949
+ python myflow.py run --with batch
988
950
  ```
989
951
  or
990
952
  ```
991
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
953
+ python myflow.py run --with kubernetes
992
954
  ```
955
+ which executes the flow on the desired system using the
956
+ requirements specified in `@resources`.
993
957
 
994
- You can also specify just the project or project branch (other values will be
995
- inferred from the current project or project branch):
996
- ```
997
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
998
- ```
999
958
 
1000
- Note that `branch` is typically one of:
1001
- - `prod`
1002
- - `user.bob`
1003
- - `test.my_experiment`
1004
- - `prod.staging`
959
+ Parameters
960
+ ----------
961
+ cpu : int, default 1
962
+ Number of CPUs required for this step.
963
+ gpu : int, optional, default None
964
+ Number of GPUs required for this step.
965
+ disk : int, optional, default None
966
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
967
+ memory : int, default 4096
968
+ Memory size (in MB) required for this step.
969
+ shared_memory : int, optional, default None
970
+ The value for the size (in MiB) of the /dev/shm volume for this step.
971
+ This parameter maps to the `--shm-size` option in Docker.
972
+ """
973
+ ...
974
+
975
+ @typing.overload
976
+ 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]]:
977
+ """
978
+ Specifies the times when the flow should be run when running on a
979
+ production scheduler.
1005
980
 
1006
981
 
1007
982
  Parameters
1008
983
  ----------
1009
- flow : Union[str, Dict[str, str]], optional, default None
1010
- Upstream flow dependency for this flow.
1011
- flows : List[Union[str, Dict[str, str]]], default []
1012
- Upstream flow dependencies for this flow.
1013
- options : Dict[str, Any], default {}
1014
- Backend-specific configuration for tuning eventing behavior.
984
+ hourly : bool, default False
985
+ Run the workflow hourly.
986
+ daily : bool, default True
987
+ Run the workflow daily.
988
+ weekly : bool, default False
989
+ Run the workflow weekly.
990
+ cron : str, optional, default None
991
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
992
+ specified by this expression.
993
+ timezone : str, optional, default None
994
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
995
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1015
996
  """
1016
997
  ...
1017
998
 
1018
999
  @typing.overload
1019
- def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1000
+ def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1020
1001
  ...
1021
1002
 
1022
- 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] = {}):
1003
+ 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):
1023
1004
  """
1024
- Specifies the flow(s) that this flow depends on.
1025
-
1026
- ```
1027
- @trigger_on_finish(flow='FooFlow')
1028
- ```
1029
- or
1030
- ```
1031
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1032
- ```
1033
- This decorator respects the @project decorator and triggers the flow
1034
- when upstream runs within the same namespace complete successfully
1035
-
1036
- Additionally, you can specify project aware upstream flow dependencies
1037
- by specifying the fully qualified project_flow_name.
1038
- ```
1039
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1040
- ```
1041
- or
1042
- ```
1043
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1044
- ```
1045
-
1046
- You can also specify just the project or project branch (other values will be
1047
- inferred from the current project or project branch):
1048
- ```
1049
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1050
- ```
1051
-
1052
- Note that `branch` is typically one of:
1053
- - `prod`
1054
- - `user.bob`
1055
- - `test.my_experiment`
1056
- - `prod.staging`
1005
+ Specifies the times when the flow should be run when running on a
1006
+ production scheduler.
1057
1007
 
1058
1008
 
1059
1009
  Parameters
1060
1010
  ----------
1061
- flow : Union[str, Dict[str, str]], optional, default None
1062
- Upstream flow dependency for this flow.
1063
- flows : List[Union[str, Dict[str, str]]], default []
1064
- Upstream flow dependencies for this flow.
1065
- options : Dict[str, Any], default {}
1066
- Backend-specific configuration for tuning eventing behavior.
1011
+ hourly : bool, default False
1012
+ Run the workflow hourly.
1013
+ daily : bool, default True
1014
+ Run the workflow daily.
1015
+ weekly : bool, default False
1016
+ Run the workflow weekly.
1017
+ cron : str, optional, default None
1018
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1019
+ specified by this expression.
1020
+ timezone : str, optional, default None
1021
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1022
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1067
1023
  """
1068
1024
  ...
1069
1025
 
@@ -1108,41 +1064,6 @@ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packag
1108
1064
  """
1109
1065
  ...
1110
1066
 
1111
- def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1112
- """
1113
- Specifies what flows belong to the same project.
1114
-
1115
- A project-specific namespace is created for all flows that
1116
- use the same `@project(name)`.
1117
-
1118
-
1119
- Parameters
1120
- ----------
1121
- name : str
1122
- Project name. Make sure that the name is unique amongst all
1123
- projects that use the same production scheduler. The name may
1124
- contain only lowercase alphanumeric characters and underscores.
1125
-
1126
- branch : Optional[str], default None
1127
- The branch to use. If not specified, the branch is set to
1128
- `user.<username>` unless `production` is set to `True`. This can
1129
- also be set on the command line using `--branch` as a top-level option.
1130
- It is an error to specify `branch` in the decorator and on the command line.
1131
-
1132
- production : bool, default False
1133
- Whether or not the branch is the production branch. This can also be set on the
1134
- command line using `--production` as a top-level option. It is an error to specify
1135
- `production` in the decorator and on the command line.
1136
- The project branch name will be:
1137
- - if `branch` is specified:
1138
- - if `production` is True: `prod.<branch>`
1139
- - if `production` is False: `test.<branch>`
1140
- - if `branch` is not specified:
1141
- - if `production` is True: `prod`
1142
- - if `production` is False: `user.<username>`
1143
- """
1144
- ...
1145
-
1146
1067
  def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None):
1147
1068
  """
1148
1069
  Allows setting external datastores to save data for the
@@ -1257,54 +1178,46 @@ def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None)
1257
1178
  """
1258
1179
  ...
1259
1180
 
1260
- @typing.overload
1261
- 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]]:
1262
- """
1263
- Specifies the Conda environment for all steps of the flow.
1264
-
1265
- Use `@conda_base` to set common libraries required by all
1266
- steps and use `@conda` to specify step-specific additions.
1267
-
1268
-
1269
- Parameters
1270
- ----------
1271
- packages : Dict[str, str], default {}
1272
- Packages to use for this flow. The key is the name of the package
1273
- and the value is the version to use.
1274
- libraries : Dict[str, str], default {}
1275
- Supported for backward compatibility. When used with packages, packages will take precedence.
1276
- python : str, optional, default None
1277
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1278
- that the version used will correspond to the version of the Python interpreter used to start the run.
1279
- disabled : bool, default False
1280
- If set to True, disables Conda.
1281
- """
1282
- ...
1283
-
1284
- @typing.overload
1285
- def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1286
- ...
1287
-
1288
- 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):
1181
+ 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]]:
1289
1182
  """
1290
- Specifies the Conda environment for all steps of the flow.
1291
-
1292
- Use `@conda_base` to set common libraries required by all
1293
- steps and use `@conda` to specify step-specific additions.
1183
+ 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)
1184
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1185
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1186
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
1187
+ starts only after all sensors finish.
1294
1188
 
1295
1189
 
1296
1190
  Parameters
1297
1191
  ----------
1298
- packages : Dict[str, str], default {}
1299
- Packages to use for this flow. The key is the name of the package
1300
- and the value is the version to use.
1301
- libraries : Dict[str, str], default {}
1302
- Supported for backward compatibility. When used with packages, packages will take precedence.
1303
- python : str, optional, default None
1304
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1305
- that the version used will correspond to the version of the Python interpreter used to start the run.
1306
- disabled : bool, default False
1307
- If set to True, disables Conda.
1192
+ timeout : int
1193
+ Time, in seconds before the task times out and fails. (Default: 3600)
1194
+ poke_interval : int
1195
+ Time in seconds that the job should wait in between each try. (Default: 60)
1196
+ mode : str
1197
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1198
+ exponential_backoff : bool
1199
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1200
+ pool : str
1201
+ the slot pool this task should run in,
1202
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1203
+ soft_fail : bool
1204
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1205
+ name : str
1206
+ Name of the sensor on Airflow
1207
+ description : str
1208
+ Description of sensor in the Airflow UI
1209
+ bucket_key : Union[str, List[str]]
1210
+ The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1211
+ When it's specified as a full s3:// url, please leave `bucket_name` as None
1212
+ bucket_name : str
1213
+ Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1214
+ When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1215
+ wildcard_match : bool
1216
+ whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1217
+ aws_conn_id : str
1218
+ a reference to the s3 connection on Airflow. (Default: None)
1219
+ verify : bool
1220
+ Whether or not to verify SSL certificates for S3 connection. (Default: None)
1308
1221
  """
1309
1222
  ...
1310
1223
 
@@ -1401,6 +1314,57 @@ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: t
1401
1314
  """
1402
1315
  ...
1403
1316
 
1317
+ @typing.overload
1318
+ 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]]:
1319
+ """
1320
+ Specifies the Conda environment for all steps of the flow.
1321
+
1322
+ Use `@conda_base` to set common libraries required by all
1323
+ steps and use `@conda` to specify step-specific additions.
1324
+
1325
+
1326
+ Parameters
1327
+ ----------
1328
+ packages : Dict[str, str], default {}
1329
+ Packages to use for this flow. The key is the name of the package
1330
+ and the value is the version to use.
1331
+ libraries : Dict[str, str], default {}
1332
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1333
+ python : str, optional, default None
1334
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1335
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1336
+ disabled : bool, default False
1337
+ If set to True, disables Conda.
1338
+ """
1339
+ ...
1340
+
1341
+ @typing.overload
1342
+ def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1343
+ ...
1344
+
1345
+ 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):
1346
+ """
1347
+ Specifies the Conda environment for all steps of the flow.
1348
+
1349
+ Use `@conda_base` to set common libraries required by all
1350
+ steps and use `@conda` to specify step-specific additions.
1351
+
1352
+
1353
+ Parameters
1354
+ ----------
1355
+ packages : Dict[str, str], default {}
1356
+ Packages to use for this flow. The key is the name of the package
1357
+ and the value is the version to use.
1358
+ libraries : Dict[str, str], default {}
1359
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1360
+ python : str, optional, default None
1361
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1362
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1363
+ disabled : bool, default False
1364
+ If set to True, disables Conda.
1365
+ """
1366
+ ...
1367
+
1404
1368
  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]]:
1405
1369
  """
1406
1370
  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.
@@ -1444,97 +1408,139 @@ def airflow_external_task_sensor(*, timeout: int, poke_interval: int, mode: str,
1444
1408
  """
1445
1409
  ...
1446
1410
 
1447
- 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]]:
1411
+ @typing.overload
1412
+ 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]]:
1448
1413
  """
1449
- 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)
1450
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1451
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1452
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
1453
- starts only after all sensors finish.
1414
+ Specifies the flow(s) that this flow depends on.
1415
+
1416
+ ```
1417
+ @trigger_on_finish(flow='FooFlow')
1418
+ ```
1419
+ or
1420
+ ```
1421
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1422
+ ```
1423
+ This decorator respects the @project decorator and triggers the flow
1424
+ when upstream runs within the same namespace complete successfully
1425
+
1426
+ Additionally, you can specify project aware upstream flow dependencies
1427
+ by specifying the fully qualified project_flow_name.
1428
+ ```
1429
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1430
+ ```
1431
+ or
1432
+ ```
1433
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1434
+ ```
1435
+
1436
+ You can also specify just the project or project branch (other values will be
1437
+ inferred from the current project or project branch):
1438
+ ```
1439
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1440
+ ```
1441
+
1442
+ Note that `branch` is typically one of:
1443
+ - `prod`
1444
+ - `user.bob`
1445
+ - `test.my_experiment`
1446
+ - `prod.staging`
1454
1447
 
1455
1448
 
1456
1449
  Parameters
1457
1450
  ----------
1458
- timeout : int
1459
- Time, in seconds before the task times out and fails. (Default: 3600)
1460
- poke_interval : int
1461
- Time in seconds that the job should wait in between each try. (Default: 60)
1462
- mode : str
1463
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1464
- exponential_backoff : bool
1465
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1466
- pool : str
1467
- the slot pool this task should run in,
1468
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1469
- soft_fail : bool
1470
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1471
- name : str
1472
- Name of the sensor on Airflow
1473
- description : str
1474
- Description of sensor in the Airflow UI
1475
- bucket_key : Union[str, List[str]]
1476
- The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1477
- When it's specified as a full s3:// url, please leave `bucket_name` as None
1478
- bucket_name : str
1479
- Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1480
- When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1481
- wildcard_match : bool
1482
- whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1483
- aws_conn_id : str
1484
- a reference to the s3 connection on Airflow. (Default: None)
1485
- verify : bool
1486
- Whether or not to verify SSL certificates for S3 connection. (Default: None)
1451
+ flow : Union[str, Dict[str, str]], optional, default None
1452
+ Upstream flow dependency for this flow.
1453
+ flows : List[Union[str, Dict[str, str]]], default []
1454
+ Upstream flow dependencies for this flow.
1455
+ options : Dict[str, Any], default {}
1456
+ Backend-specific configuration for tuning eventing behavior.
1487
1457
  """
1488
1458
  ...
1489
1459
 
1490
1460
  @typing.overload
1491
- 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]]:
1461
+ def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1462
+ ...
1463
+
1464
+ 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] = {}):
1492
1465
  """
1493
- Specifies the times when the flow should be run when running on a
1494
- production scheduler.
1466
+ Specifies the flow(s) that this flow depends on.
1467
+
1468
+ ```
1469
+ @trigger_on_finish(flow='FooFlow')
1470
+ ```
1471
+ or
1472
+ ```
1473
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1474
+ ```
1475
+ This decorator respects the @project decorator and triggers the flow
1476
+ when upstream runs within the same namespace complete successfully
1477
+
1478
+ Additionally, you can specify project aware upstream flow dependencies
1479
+ by specifying the fully qualified project_flow_name.
1480
+ ```
1481
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1482
+ ```
1483
+ or
1484
+ ```
1485
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1486
+ ```
1487
+
1488
+ You can also specify just the project or project branch (other values will be
1489
+ inferred from the current project or project branch):
1490
+ ```
1491
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1492
+ ```
1493
+
1494
+ Note that `branch` is typically one of:
1495
+ - `prod`
1496
+ - `user.bob`
1497
+ - `test.my_experiment`
1498
+ - `prod.staging`
1495
1499
 
1496
1500
 
1497
1501
  Parameters
1498
1502
  ----------
1499
- hourly : bool, default False
1500
- Run the workflow hourly.
1501
- daily : bool, default True
1502
- Run the workflow daily.
1503
- weekly : bool, default False
1504
- Run the workflow weekly.
1505
- cron : str, optional, default None
1506
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1507
- specified by this expression.
1508
- timezone : str, optional, default None
1509
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1510
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1503
+ flow : Union[str, Dict[str, str]], optional, default None
1504
+ Upstream flow dependency for this flow.
1505
+ flows : List[Union[str, Dict[str, str]]], default []
1506
+ Upstream flow dependencies for this flow.
1507
+ options : Dict[str, Any], default {}
1508
+ Backend-specific configuration for tuning eventing behavior.
1511
1509
  """
1512
1510
  ...
1513
1511
 
1514
- @typing.overload
1515
- def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1516
- ...
1517
-
1518
- def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None):
1512
+ def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1519
1513
  """
1520
- Specifies the times when the flow should be run when running on a
1521
- production scheduler.
1514
+ Specifies what flows belong to the same project.
1515
+
1516
+ A project-specific namespace is created for all flows that
1517
+ use the same `@project(name)`.
1522
1518
 
1523
1519
 
1524
1520
  Parameters
1525
1521
  ----------
1526
- hourly : bool, default False
1527
- Run the workflow hourly.
1528
- daily : bool, default True
1529
- Run the workflow daily.
1530
- weekly : bool, default False
1531
- Run the workflow weekly.
1532
- cron : str, optional, default None
1533
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1534
- specified by this expression.
1535
- timezone : str, optional, default None
1536
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1537
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1522
+ name : str
1523
+ Project name. Make sure that the name is unique amongst all
1524
+ projects that use the same production scheduler. The name may
1525
+ contain only lowercase alphanumeric characters and underscores.
1526
+
1527
+ branch : Optional[str], default None
1528
+ The branch to use. If not specified, the branch is set to
1529
+ `user.<username>` unless `production` is set to `True`. This can
1530
+ also be set on the command line using `--branch` as a top-level option.
1531
+ It is an error to specify `branch` in the decorator and on the command line.
1532
+
1533
+ production : bool, default False
1534
+ Whether or not the branch is the production branch. This can also be set on the
1535
+ command line using `--production` as a top-level option. It is an error to specify
1536
+ `production` in the decorator and on the command line.
1537
+ The project branch name will be:
1538
+ - if `branch` is specified:
1539
+ - if `production` is True: `prod.<branch>`
1540
+ - if `production` is False: `test.<branch>`
1541
+ - if `branch` is not specified:
1542
+ - if `production` is True: `prod`
1543
+ - if `production` is False: `user.<username>`
1538
1544
  """
1539
1545
  ...
1540
1546