ob-metaflow-stubs 6.0.3.151__py2.py3-none-any.whl → 6.0.3.152__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 (198) hide show
  1. metaflow-stubs/__init__.pyi +698 -683
  2. metaflow-stubs/cards.pyi +2 -2
  3. metaflow-stubs/cli.pyi +2 -2
  4. metaflow-stubs/cli_components/__init__.pyi +2 -2
  5. metaflow-stubs/cli_components/utils.pyi +2 -2
  6. metaflow-stubs/client/__init__.pyi +2 -2
  7. metaflow-stubs/client/core.pyi +5 -5
  8. metaflow-stubs/client/filecache.pyi +2 -2
  9. metaflow-stubs/events.pyi +2 -2
  10. metaflow-stubs/exception.pyi +2 -2
  11. metaflow-stubs/flowspec.pyi +4 -4
  12. metaflow-stubs/generated_for.txt +1 -1
  13. metaflow-stubs/includefile.pyi +3 -3
  14. metaflow-stubs/info_file.pyi +2 -2
  15. metaflow-stubs/metadata_provider/__init__.pyi +2 -2
  16. metaflow-stubs/metadata_provider/heartbeat.pyi +2 -2
  17. metaflow-stubs/metadata_provider/metadata.pyi +2 -2
  18. metaflow-stubs/metadata_provider/util.pyi +2 -2
  19. metaflow-stubs/metaflow_config.pyi +2 -2
  20. metaflow-stubs/metaflow_current.pyi +97 -97
  21. metaflow-stubs/mf_extensions/__init__.pyi +2 -2
  22. metaflow-stubs/mf_extensions/obcheckpoint/__init__.pyi +2 -2
  23. metaflow-stubs/mf_extensions/obcheckpoint/plugins/__init__.pyi +2 -2
  24. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/__init__.pyi +2 -2
  25. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/__init__.pyi +2 -2
  26. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/async_cards.pyi +3 -3
  27. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/deco_injection_mixin.pyi +2 -2
  28. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/extra_components.pyi +2 -2
  29. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/__init__.pyi +2 -2
  30. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/__init__.pyi +2 -2
  31. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/checkpoint_lister.pyi +4 -4
  32. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/lineage_card.pyi +2 -2
  33. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/checkpoint_storage.pyi +4 -4
  34. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/constructors.pyi +2 -2
  35. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/core.pyi +3 -3
  36. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/decorator.pyi +3 -3
  37. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/exceptions.pyi +2 -2
  38. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/final_api.pyi +2 -2
  39. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/lineage.pyi +2 -2
  40. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/__init__.pyi +2 -2
  41. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/core.pyi +3 -3
  42. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/exceptions.pyi +2 -2
  43. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/task_utils.pyi +2 -2
  44. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/utils.pyi +2 -2
  45. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastructures.pyi +2 -2
  46. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/exceptions.pyi +2 -2
  47. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/__init__.pyi +2 -2
  48. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/decorator.pyi +3 -3
  49. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/__init__.pyi +2 -2
  50. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/core.pyi +3 -3
  51. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/exceptions.pyi +2 -2
  52. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/model_storage.pyi +3 -3
  53. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/__init__.pyi +2 -2
  54. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/flowspec_utils.pyi +2 -2
  55. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/general.pyi +2 -2
  56. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/identity_utils.pyi +2 -2
  57. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/__init__.pyi +2 -2
  58. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/base.pyi +2 -2
  59. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/tar.pyi +2 -2
  60. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/tar_utils.pyi +2 -2
  61. metaflow-stubs/mf_extensions/outerbounds/__init__.pyi +2 -2
  62. metaflow-stubs/mf_extensions/outerbounds/plugins/__init__.pyi +2 -2
  63. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/__init__.pyi +2 -2
  64. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/injector.pyi +2 -2
  65. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/__init__.pyi +2 -2
  66. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/ollama.pyi +2 -2
  67. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/__init__.pyi +2 -2
  68. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/snowflake.pyi +2 -2
  69. metaflow-stubs/mf_extensions/outerbounds/profilers/__init__.pyi +2 -2
  70. metaflow-stubs/mf_extensions/outerbounds/profilers/gpu.pyi +2 -2
  71. metaflow-stubs/mf_extensions/outerbounds/remote_config.pyi +2 -2
  72. metaflow-stubs/mf_extensions/outerbounds/toplevel/__init__.pyi +2 -2
  73. metaflow-stubs/mf_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.pyi +2 -2
  74. metaflow-stubs/multicore_utils.pyi +2 -2
  75. metaflow-stubs/parameters.pyi +3 -3
  76. metaflow-stubs/plugins/__init__.pyi +13 -13
  77. metaflow-stubs/plugins/airflow/__init__.pyi +2 -2
  78. metaflow-stubs/plugins/airflow/airflow_utils.pyi +2 -2
  79. metaflow-stubs/plugins/airflow/exception.pyi +2 -2
  80. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +2 -2
  81. metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi +2 -2
  82. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +2 -2
  83. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +2 -2
  84. metaflow-stubs/plugins/argo/__init__.pyi +2 -2
  85. metaflow-stubs/plugins/argo/argo_client.pyi +2 -2
  86. metaflow-stubs/plugins/argo/argo_events.pyi +2 -2
  87. metaflow-stubs/plugins/argo/argo_workflows.pyi +2 -2
  88. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +2 -2
  89. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +3 -3
  90. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +3 -3
  91. metaflow-stubs/plugins/aws/__init__.pyi +2 -2
  92. metaflow-stubs/plugins/aws/aws_client.pyi +2 -2
  93. metaflow-stubs/plugins/aws/aws_utils.pyi +2 -2
  94. metaflow-stubs/plugins/aws/batch/__init__.pyi +2 -2
  95. metaflow-stubs/plugins/aws/batch/batch.pyi +2 -2
  96. metaflow-stubs/plugins/aws/batch/batch_client.pyi +2 -2
  97. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +2 -2
  98. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +2 -2
  99. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +3 -3
  100. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +2 -2
  101. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +2 -2
  102. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +2 -2
  103. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +2 -2
  104. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +2 -2
  105. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer.pyi +3 -3
  106. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer_objects.pyi +3 -3
  107. metaflow-stubs/plugins/azure/__init__.pyi +2 -2
  108. metaflow-stubs/plugins/azure/azure_credential.pyi +2 -2
  109. metaflow-stubs/plugins/azure/azure_exceptions.pyi +2 -2
  110. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +3 -3
  111. metaflow-stubs/plugins/azure/azure_utils.pyi +2 -2
  112. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +2 -2
  113. metaflow-stubs/plugins/azure/includefile_support.pyi +2 -2
  114. metaflow-stubs/plugins/cards/__init__.pyi +2 -2
  115. metaflow-stubs/plugins/cards/card_client.pyi +2 -2
  116. metaflow-stubs/plugins/cards/card_creator.pyi +2 -2
  117. metaflow-stubs/plugins/cards/card_datastore.pyi +2 -2
  118. metaflow-stubs/plugins/cards/card_decorator.pyi +2 -2
  119. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +2 -2
  120. metaflow-stubs/plugins/cards/card_modules/basic.pyi +2 -2
  121. metaflow-stubs/plugins/cards/card_modules/card.pyi +2 -2
  122. metaflow-stubs/plugins/cards/card_modules/components.pyi +3 -3
  123. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +2 -2
  124. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +2 -2
  125. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +2 -2
  126. metaflow-stubs/plugins/cards/card_resolver.pyi +2 -2
  127. metaflow-stubs/plugins/cards/component_serializer.pyi +2 -2
  128. metaflow-stubs/plugins/cards/exception.pyi +2 -2
  129. metaflow-stubs/plugins/catch_decorator.pyi +2 -2
  130. metaflow-stubs/plugins/datatools/__init__.pyi +2 -2
  131. metaflow-stubs/plugins/datatools/local.pyi +2 -2
  132. metaflow-stubs/plugins/datatools/s3/__init__.pyi +2 -2
  133. metaflow-stubs/plugins/datatools/s3/s3.pyi +2 -2
  134. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +2 -2
  135. metaflow-stubs/plugins/datatools/s3/s3util.pyi +2 -2
  136. metaflow-stubs/plugins/debug_logger.pyi +2 -2
  137. metaflow-stubs/plugins/debug_monitor.pyi +2 -2
  138. metaflow-stubs/plugins/environment_decorator.pyi +2 -2
  139. metaflow-stubs/plugins/events_decorator.pyi +2 -2
  140. metaflow-stubs/plugins/frameworks/__init__.pyi +2 -2
  141. metaflow-stubs/plugins/frameworks/pytorch.pyi +2 -2
  142. metaflow-stubs/plugins/gcp/__init__.pyi +2 -2
  143. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +3 -3
  144. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +2 -2
  145. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +2 -2
  146. metaflow-stubs/plugins/gcp/gs_utils.pyi +2 -2
  147. metaflow-stubs/plugins/gcp/includefile_support.pyi +2 -2
  148. metaflow-stubs/plugins/kubernetes/__init__.pyi +2 -2
  149. metaflow-stubs/plugins/kubernetes/kube_utils.pyi +2 -2
  150. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +2 -2
  151. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +2 -2
  152. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +2 -2
  153. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +2 -2
  154. metaflow-stubs/plugins/kubernetes/spot_monitor_sidecar.pyi +2 -2
  155. metaflow-stubs/plugins/ollama/__init__.pyi +2 -2
  156. metaflow-stubs/plugins/parallel_decorator.pyi +2 -2
  157. metaflow-stubs/plugins/perimeters.pyi +2 -2
  158. metaflow-stubs/plugins/project_decorator.pyi +2 -2
  159. metaflow-stubs/plugins/pypi/__init__.pyi +2 -2
  160. metaflow-stubs/plugins/pypi/conda_decorator.pyi +2 -2
  161. metaflow-stubs/plugins/pypi/conda_environment.pyi +3 -3
  162. metaflow-stubs/plugins/pypi/parsers.pyi +2 -2
  163. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +2 -2
  164. metaflow-stubs/plugins/pypi/pypi_environment.pyi +2 -2
  165. metaflow-stubs/plugins/pypi/utils.pyi +2 -2
  166. metaflow-stubs/plugins/resources_decorator.pyi +2 -2
  167. metaflow-stubs/plugins/retry_decorator.pyi +2 -2
  168. metaflow-stubs/plugins/secrets/__init__.pyi +3 -3
  169. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +3 -3
  170. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +2 -2
  171. metaflow-stubs/plugins/snowflake/__init__.pyi +2 -2
  172. metaflow-stubs/plugins/storage_executor.pyi +2 -2
  173. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +2 -2
  174. metaflow-stubs/plugins/timeout_decorator.pyi +2 -2
  175. metaflow-stubs/profilers/__init__.pyi +2 -2
  176. metaflow-stubs/pylint_wrapper.pyi +2 -2
  177. metaflow-stubs/runner/__init__.pyi +2 -2
  178. metaflow-stubs/runner/deployer.pyi +29 -29
  179. metaflow-stubs/runner/deployer_impl.pyi +3 -3
  180. metaflow-stubs/runner/metaflow_runner.pyi +3 -3
  181. metaflow-stubs/runner/nbdeploy.pyi +2 -2
  182. metaflow-stubs/runner/nbrun.pyi +2 -2
  183. metaflow-stubs/runner/subprocess_manager.pyi +2 -2
  184. metaflow-stubs/runner/utils.pyi +4 -4
  185. metaflow-stubs/system/__init__.pyi +2 -2
  186. metaflow-stubs/system/system_logger.pyi +3 -3
  187. metaflow-stubs/system/system_monitor.pyi +2 -2
  188. metaflow-stubs/tagging_util.pyi +2 -2
  189. metaflow-stubs/tuple_util.pyi +2 -2
  190. metaflow-stubs/user_configs/__init__.pyi +2 -2
  191. metaflow-stubs/user_configs/config_decorators.pyi +7 -7
  192. metaflow-stubs/user_configs/config_options.pyi +3 -3
  193. metaflow-stubs/user_configs/config_parameters.pyi +8 -8
  194. {ob_metaflow_stubs-6.0.3.151.dist-info → ob_metaflow_stubs-6.0.3.152.dist-info}/METADATA +1 -1
  195. ob_metaflow_stubs-6.0.3.152.dist-info/RECORD +198 -0
  196. ob_metaflow_stubs-6.0.3.151.dist-info/RECORD +0 -198
  197. {ob_metaflow_stubs-6.0.3.151.dist-info → ob_metaflow_stubs-6.0.3.152.dist-info}/WHEEL +0 -0
  198. {ob_metaflow_stubs-6.0.3.151.dist-info → ob_metaflow_stubs-6.0.3.152.dist-info}/top_level.txt +0 -0
@@ -1,15 +1,15 @@
1
1
  ######################################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
- # MF version: 2.15.6.1+obcheckpoint(0.1.9);ob(v1) #
4
- # Generated on 2025-03-17T20:51:51.819256 #
3
+ # MF version: 2.15.7.1+obcheckpoint(0.1.9);ob(v1) #
4
+ # Generated on 2025-03-24T18:51:07.966816 #
5
5
  ######################################################################################################
6
6
 
7
7
  from __future__ import annotations
8
8
 
9
9
  import typing
10
10
  if typing.TYPE_CHECKING:
11
- import datetime
12
11
  import typing
12
+ import datetime
13
13
  FlowSpecDerived = typing.TypeVar("FlowSpecDerived", bound="FlowSpec", contravariant=False, covariant=False)
14
14
  StepFlag = typing.NewType("StepFlag", bool)
15
15
 
@@ -43,9 +43,9 @@ from . import plugins as plugins
43
43
  from .mf_extensions.outerbounds.toplevel.global_aliases_for_metaflow_package import S3 as S3
44
44
  from . import includefile as includefile
45
45
  from .includefile import IncludeFile as IncludeFile
46
- from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
47
46
  from .plugins.pypi.parsers import conda_environment_yml_parser as conda_environment_yml_parser
48
47
  from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
48
+ from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
49
49
  from . import client as client
50
50
  from .client.core import namespace as namespace
51
51
  from .client.core import get_namespace as get_namespace
@@ -149,273 +149,446 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
149
149
  """
150
150
  ...
151
151
 
152
- def ollama(*, models: "list[Ollama]", backend: 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]]]:
152
+ @typing.overload
153
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
153
154
  """
154
- This decorator is used to run Ollama APIs as Metaflow task sidecars.
155
-
156
- User code call
157
- -----------
158
- @ollama(
159
- models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
160
- backend='local'
161
- )
162
-
163
- Valid backend options
164
- ---------------------
165
- - 'local': Run as a separate process on the local task machine.
166
- - (TODO) 'managed': Outerbounds hosts and selects compute provider.
167
- - (TODO) 'remote': Spin up separate instance to serve Ollama models.
168
-
169
- Valid model options
170
- ----------------
171
- - 'llama3.2'
172
- - 'llama3.3'
173
- - any model here https://ollama.com/search
174
-
175
-
176
- Parameters
177
- ----------
178
- models: list[Ollama]
179
- List of Ollama containers running models in sidecars.
180
- backend: str
181
- Determines where and how to run the Ollama process.
155
+ Internal decorator to support Fast bakery
182
156
  """
183
157
  ...
184
158
 
185
159
  @typing.overload
186
- 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]]]:
160
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
161
+ ...
162
+
163
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
187
164
  """
188
- Specifies environment variables to be set prior to the execution of a step.
165
+ Internal decorator to support Fast bakery
166
+ """
167
+ ...
168
+
169
+ @typing.overload
170
+ 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]]]:
171
+ """
172
+ Enables loading / saving of models within a step.
173
+
189
174
 
190
175
 
191
176
  Parameters
192
177
  ----------
193
- vars : Dict[str, str], default {}
194
- Dictionary of environment variables to set.
178
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
179
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
180
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
181
+ - `current.checkpoint`
182
+ - `current.model`
183
+ - `current.huggingface_hub`
184
+
185
+ 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
186
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
187
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
188
+
189
+ temp_dir_root : str, default: None
190
+ The root directory under which `current.model.loaded` will store loaded models
195
191
  """
196
192
  ...
197
193
 
198
194
  @typing.overload
199
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
195
+ def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
200
196
  ...
201
197
 
202
198
  @typing.overload
203
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
199
+ def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
204
200
  ...
205
201
 
206
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
202
+ 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):
207
203
  """
208
- Specifies environment variables to be set prior to the execution of a step.
204
+ Enables loading / saving of models within a step.
205
+
209
206
 
210
207
 
211
208
  Parameters
212
209
  ----------
213
- vars : Dict[str, str], default {}
214
- Dictionary of environment variables to set.
210
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
211
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
212
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
213
+ - `current.checkpoint`
214
+ - `current.model`
215
+ - `current.huggingface_hub`
216
+
217
+ 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
218
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
219
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
220
+
221
+ temp_dir_root : str, default: None
222
+ The root directory under which `current.model.loaded` will store loaded models
215
223
  """
216
224
  ...
217
225
 
218
226
  @typing.overload
219
- 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]]]:
227
+ 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]]]:
220
228
  """
221
- Specifies the number of times the task corresponding
222
- to a step needs to be retried.
223
-
224
- This decorator is useful for handling transient errors, such as networking issues.
225
- If your task contains operations that can't be retried safely, e.g. database updates,
226
- it is advisable to annotate it with `@retry(times=0)`.
229
+ Specifies the Conda environment for the step.
227
230
 
228
- This can be used in conjunction with the `@catch` decorator. The `@catch`
229
- decorator will execute a no-op task after all retries have been exhausted,
230
- ensuring that the flow execution can continue.
231
+ Information in this decorator will augment any
232
+ attributes set in the `@conda_base` flow-level decorator. Hence,
233
+ you can use `@conda_base` to set packages required by all
234
+ steps and use `@conda` to specify step-specific overrides.
231
235
 
232
236
 
233
237
  Parameters
234
238
  ----------
235
- times : int, default 3
236
- Number of times to retry this task.
237
- minutes_between_retries : int, default 2
238
- Number of minutes between retries.
239
+ packages : Dict[str, str], default {}
240
+ Packages to use for this step. The key is the name of the package
241
+ and the value is the version to use.
242
+ libraries : Dict[str, str], default {}
243
+ Supported for backward compatibility. When used with packages, packages will take precedence.
244
+ python : str, optional, default None
245
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
246
+ that the version used will correspond to the version of the Python interpreter used to start the run.
247
+ disabled : bool, default False
248
+ If set to True, disables @conda.
239
249
  """
240
250
  ...
241
251
 
242
252
  @typing.overload
243
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
253
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
244
254
  ...
245
255
 
246
256
  @typing.overload
247
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
257
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
248
258
  ...
249
259
 
250
- 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):
260
+ 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):
251
261
  """
252
- Specifies the number of times the task corresponding
253
- to a step needs to be retried.
254
-
255
- This decorator is useful for handling transient errors, such as networking issues.
256
- If your task contains operations that can't be retried safely, e.g. database updates,
257
- it is advisable to annotate it with `@retry(times=0)`.
262
+ Specifies the Conda environment for the step.
258
263
 
259
- This can be used in conjunction with the `@catch` decorator. The `@catch`
260
- decorator will execute a no-op task after all retries have been exhausted,
261
- ensuring that the flow execution can continue.
264
+ Information in this decorator will augment any
265
+ attributes set in the `@conda_base` flow-level decorator. Hence,
266
+ you can use `@conda_base` to set packages required by all
267
+ steps and use `@conda` to specify step-specific overrides.
262
268
 
263
269
 
264
270
  Parameters
265
271
  ----------
266
- times : int, default 3
267
- Number of times to retry this task.
268
- minutes_between_retries : int, default 2
269
- Number of minutes between retries.
272
+ packages : Dict[str, str], default {}
273
+ Packages to use for this step. The key is the name of the package
274
+ and the value is the version to use.
275
+ libraries : Dict[str, str], default {}
276
+ Supported for backward compatibility. When used with packages, packages will take precedence.
277
+ python : str, optional, default None
278
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
279
+ that the version used will correspond to the version of the Python interpreter used to start the run.
280
+ disabled : bool, default False
281
+ If set to True, disables @conda.
270
282
  """
271
283
  ...
272
284
 
273
285
  @typing.overload
274
- 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]]]:
286
+ 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]]]:
275
287
  """
276
- Specifies secrets to be retrieved and injected as environment variables prior to
277
- the execution of a step.
288
+ Specifies that the step will success under all circumstances.
289
+
290
+ The decorator will create an optional artifact, specified by `var`, which
291
+ contains the exception raised. You can use it to detect the presence
292
+ of errors, indicating that all happy-path artifacts produced by the step
293
+ are missing.
278
294
 
279
295
 
280
296
  Parameters
281
297
  ----------
282
- sources : List[Union[str, Dict[str, Any]]], default: []
283
- List of secret specs, defining how the secrets are to be retrieved
298
+ var : str, optional, default None
299
+ Name of the artifact in which to store the caught exception.
300
+ If not specified, the exception is not stored.
301
+ print_exception : bool, default True
302
+ Determines whether or not the exception is printed to
303
+ stdout when caught.
284
304
  """
285
305
  ...
286
306
 
287
307
  @typing.overload
288
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
308
+ def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
289
309
  ...
290
310
 
291
311
  @typing.overload
292
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
312
+ def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
293
313
  ...
294
314
 
295
- 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]]] = []):
315
+ 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):
296
316
  """
297
- Specifies secrets to be retrieved and injected as environment variables prior to
298
- the execution of a step.
317
+ Specifies that the step will success under all circumstances.
318
+
319
+ The decorator will create an optional artifact, specified by `var`, which
320
+ contains the exception raised. You can use it to detect the presence
321
+ of errors, indicating that all happy-path artifacts produced by the step
322
+ are missing.
299
323
 
300
324
 
301
325
  Parameters
302
326
  ----------
303
- sources : List[Union[str, Dict[str, Any]]], default: []
304
- List of secret specs, defining how the secrets are to be retrieved
327
+ var : str, optional, default None
328
+ Name of the artifact in which to store the caught exception.
329
+ If not specified, the exception is not stored.
330
+ print_exception : bool, default True
331
+ Determines whether or not the exception is printed to
332
+ stdout when caught.
333
+ """
334
+ ...
335
+
336
+ 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]]]:
337
+ """
338
+ Decorator that helps cache, version and store models/datasets from huggingface hub.
339
+
340
+
341
+ Parameters
342
+ ----------
343
+ temp_dir_root : str, optional
344
+ The root directory that will hold the temporary directory where objects will be downloaded.
345
+
346
+ load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
347
+ The list of repos (models/datasets) to load.
348
+
349
+ Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
350
+
351
+ - If repo (model/dataset) is not found in the datastore:
352
+ - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
353
+ - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
354
+ - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
355
+
356
+ - If repo is found in the datastore:
357
+ - Loads it directly from datastore to local path (can be temporary directory or specified path)
305
358
  """
306
359
  ...
307
360
 
308
361
  @typing.overload
309
- 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]]]:
362
+ def pypi(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
310
363
  """
311
- Enables checkpointing for a step.
364
+ Specifies the PyPI packages for the step.
312
365
 
366
+ Information in this decorator will augment any
367
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
368
+ you can use `@pypi_base` to set packages required by all
369
+ steps and use `@pypi` to specify step-specific overrides.
313
370
 
314
371
 
315
372
  Parameters
316
373
  ----------
317
- load_policy : str, default: "fresh"
318
- The policy for loading the checkpoint. The following policies are supported:
319
- - "eager": Loads the the latest available checkpoint within the namespace.
320
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
321
- will be loaded at the start of the task.
322
- - "none": Do not load any checkpoint
323
- - "fresh": Loads the lastest checkpoint created within the running Task.
324
- This mode helps loading checkpoints across various retry attempts of the same task.
325
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
326
- created within the task will be loaded when the task is retries execution on failure.
327
-
328
- temp_dir_root : str, default: None
329
- The root directory under which `current.checkpoint.directory` will be created.
374
+ packages : Dict[str, str], default: {}
375
+ Packages to use for this step. The key is the name of the package
376
+ and the value is the version to use.
377
+ python : str, optional, default: None
378
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
379
+ that the version used will correspond to the version of the Python interpreter used to start the run.
330
380
  """
331
381
  ...
332
382
 
333
383
  @typing.overload
334
- def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
384
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
335
385
  ...
336
386
 
337
387
  @typing.overload
338
- def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
388
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
339
389
  ...
340
390
 
341
- 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):
391
+ def pypi(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
342
392
  """
343
- Enables checkpointing for a step.
344
-
345
-
393
+ Specifies the PyPI packages for the step.
394
+
395
+ Information in this decorator will augment any
396
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
397
+ you can use `@pypi_base` to set packages required by all
398
+ steps and use `@pypi` to specify step-specific overrides.
399
+
346
400
 
347
401
  Parameters
348
402
  ----------
349
- load_policy : str, default: "fresh"
350
- The policy for loading the checkpoint. The following policies are supported:
351
- - "eager": Loads the the latest available checkpoint within the namespace.
352
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
353
- will be loaded at the start of the task.
354
- - "none": Do not load any checkpoint
355
- - "fresh": Loads the lastest checkpoint created within the running Task.
356
- This mode helps loading checkpoints across various retry attempts of the same task.
357
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
358
- created within the task will be loaded when the task is retries execution on failure.
403
+ packages : Dict[str, str], default: {}
404
+ Packages to use for this step. The key is the name of the package
405
+ and the value is the version to use.
406
+ python : str, optional, default: None
407
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
408
+ that the version used will correspond to the version of the Python interpreter used to start the run.
409
+ """
410
+ ...
411
+
412
+ 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]]]:
413
+ """
414
+ Specifies that this step should execute on DGX cloud.
359
415
 
360
- temp_dir_root : str, default: None
361
- The root directory under which `current.checkpoint.directory` will be created.
416
+
417
+ Parameters
418
+ ----------
419
+ gpu : int
420
+ Number of GPUs to use.
421
+ gpu_type : str
422
+ Type of Nvidia GPU to use.
423
+ queue_timeout : int
424
+ Time to keep the job in NVCF's queue.
362
425
  """
363
426
  ...
364
427
 
365
- @typing.overload
366
- 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]]]:
428
+ def nim(*, models: "list[NIM]", backend: 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]]]:
367
429
  """
368
- Enables loading / saving of models within a step.
430
+ This decorator is used to run NIM containers in Metaflow tasks as sidecars.
431
+
432
+ User code call
433
+ -----------
434
+ @nim(
435
+ models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
436
+ backend='managed'
437
+ )
438
+
439
+ Valid backend options
440
+ ---------------------
441
+ - 'managed': Outerbounds selects a compute provider based on the model.
369
442
 
443
+ Valid model options
444
+ ----------------
445
+ - 'meta/llama3-8b-instruct': 8B parameter model
446
+ - 'meta/llama3-70b-instruct': 70B parameter model
447
+ - any model here: https://nvcf.ngc.nvidia.com/functions?filter=nvidia-functions
370
448
 
371
449
 
372
450
  Parameters
373
451
  ----------
374
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
375
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
376
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
377
- - `current.checkpoint`
378
- - `current.model`
379
- - `current.huggingface_hub`
452
+ models: list[NIM]
453
+ List of NIM containers running models in sidecars.
454
+ backend: str
455
+ Compute provider to run the NIM container.
456
+ queue_timeout : int
457
+ Time to keep the job in NVCF's queue.
458
+ """
459
+ ...
460
+
461
+ @typing.overload
462
+ 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]]]:
463
+ """
464
+ Specifies the number of times the task corresponding
465
+ to a step needs to be retried.
380
466
 
381
- 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
382
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
383
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
467
+ This decorator is useful for handling transient errors, such as networking issues.
468
+ If your task contains operations that can't be retried safely, e.g. database updates,
469
+ it is advisable to annotate it with `@retry(times=0)`.
384
470
 
385
- temp_dir_root : str, default: None
386
- The root directory under which `current.model.loaded` will store loaded models
471
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
472
+ decorator will execute a no-op task after all retries have been exhausted,
473
+ ensuring that the flow execution can continue.
474
+
475
+
476
+ Parameters
477
+ ----------
478
+ times : int, default 3
479
+ Number of times to retry this task.
480
+ minutes_between_retries : int, default 2
481
+ Number of minutes between retries.
387
482
  """
388
483
  ...
389
484
 
390
485
  @typing.overload
391
- def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
486
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
392
487
  ...
393
488
 
394
489
  @typing.overload
395
- def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
490
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
396
491
  ...
397
492
 
398
- 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):
493
+ 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):
399
494
  """
400
- Enables loading / saving of models within a step.
495
+ Specifies the number of times the task corresponding
496
+ to a step needs to be retried.
497
+
498
+ This decorator is useful for handling transient errors, such as networking issues.
499
+ If your task contains operations that can't be retried safely, e.g. database updates,
500
+ it is advisable to annotate it with `@retry(times=0)`.
401
501
 
502
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
503
+ decorator will execute a no-op task after all retries have been exhausted,
504
+ ensuring that the flow execution can continue.
402
505
 
403
506
 
404
507
  Parameters
405
508
  ----------
406
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
407
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
408
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
409
- - `current.checkpoint`
410
- - `current.model`
411
- - `current.huggingface_hub`
509
+ times : int, default 3
510
+ Number of times to retry this task.
511
+ minutes_between_retries : int, default 2
512
+ Number of minutes between retries.
513
+ """
514
+ ...
515
+
516
+ @typing.overload
517
+ 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]]]:
518
+ """
519
+ Specifies the resources needed when executing this step.
412
520
 
413
- 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
414
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
415
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
521
+ Use `@resources` to specify the resource requirements
522
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
416
523
 
417
- temp_dir_root : str, default: None
418
- The root directory under which `current.model.loaded` will store loaded models
524
+ You can choose the compute layer on the command line by executing e.g.
525
+ ```
526
+ python myflow.py run --with batch
527
+ ```
528
+ or
529
+ ```
530
+ python myflow.py run --with kubernetes
531
+ ```
532
+ which executes the flow on the desired system using the
533
+ requirements specified in `@resources`.
534
+
535
+
536
+ Parameters
537
+ ----------
538
+ cpu : int, default 1
539
+ Number of CPUs required for this step.
540
+ gpu : int, optional, default None
541
+ Number of GPUs required for this step.
542
+ disk : int, optional, default None
543
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
544
+ memory : int, default 4096
545
+ Memory size (in MB) required for this step.
546
+ shared_memory : int, optional, default None
547
+ The value for the size (in MiB) of the /dev/shm volume for this step.
548
+ This parameter maps to the `--shm-size` option in Docker.
549
+ """
550
+ ...
551
+
552
+ @typing.overload
553
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
554
+ ...
555
+
556
+ @typing.overload
557
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
558
+ ...
559
+
560
+ 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):
561
+ """
562
+ Specifies the resources needed when executing this step.
563
+
564
+ Use `@resources` to specify the resource requirements
565
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
566
+
567
+ You can choose the compute layer on the command line by executing e.g.
568
+ ```
569
+ python myflow.py run --with batch
570
+ ```
571
+ or
572
+ ```
573
+ python myflow.py run --with kubernetes
574
+ ```
575
+ which executes the flow on the desired system using the
576
+ requirements specified in `@resources`.
577
+
578
+
579
+ Parameters
580
+ ----------
581
+ cpu : int, default 1
582
+ Number of CPUs required for this step.
583
+ gpu : int, optional, default None
584
+ Number of GPUs required for this step.
585
+ disk : int, optional, default None
586
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
587
+ memory : int, default 4096
588
+ Memory size (in MB) required for this step.
589
+ shared_memory : int, optional, default None
590
+ The value for the size (in MiB) of the /dev/shm volume for this step.
591
+ This parameter maps to the `--shm-size` option in Docker.
419
592
  """
420
593
  ...
421
594
 
@@ -497,102 +670,40 @@ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: ty
497
670
  ...
498
671
 
499
672
  @typing.overload
500
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
501
- """
502
- Decorator prototype for all step decorators. This function gets specialized
503
- and imported for all decorators types by _import_plugin_decorators().
504
- """
505
- ...
506
-
507
- @typing.overload
508
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
509
- ...
510
-
511
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
673
+ 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]]]:
512
674
  """
513
- Decorator prototype for all step decorators. This function gets specialized
514
- and imported for all decorators types by _import_plugin_decorators().
675
+ Specifies a timeout for your step.
676
+
677
+ This decorator is useful if this step may hang indefinitely.
678
+
679
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
680
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
681
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
682
+
683
+ Note that all the values specified in parameters are added together so if you specify
684
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
685
+
686
+
687
+ Parameters
688
+ ----------
689
+ seconds : int, default 0
690
+ Number of seconds to wait prior to timing out.
691
+ minutes : int, default 0
692
+ Number of minutes to wait prior to timing out.
693
+ hours : int, default 0
694
+ Number of hours to wait prior to timing out.
515
695
  """
516
696
  ...
517
697
 
518
698
  @typing.overload
519
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
520
- """
521
- Internal decorator to support Fast bakery
522
- """
699
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
523
700
  ...
524
701
 
525
702
  @typing.overload
526
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
703
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
527
704
  ...
528
705
 
529
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
530
- """
531
- Internal decorator to support Fast bakery
532
- """
533
- ...
534
-
535
- @typing.overload
536
- 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]]]:
537
- """
538
- Specifies the Conda environment for the step.
539
-
540
- Information in this decorator will augment any
541
- attributes set in the `@conda_base` flow-level decorator. Hence,
542
- you can use `@conda_base` to set packages required by all
543
- steps and use `@conda` to specify step-specific overrides.
544
-
545
-
546
- Parameters
547
- ----------
548
- packages : Dict[str, str], default {}
549
- Packages to use for this step. The key is the name of the package
550
- and the value is the version to use.
551
- libraries : Dict[str, str], default {}
552
- Supported for backward compatibility. When used with packages, packages will take precedence.
553
- python : str, optional, default None
554
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
555
- that the version used will correspond to the version of the Python interpreter used to start the run.
556
- disabled : bool, default False
557
- If set to True, disables @conda.
558
- """
559
- ...
560
-
561
- @typing.overload
562
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
563
- ...
564
-
565
- @typing.overload
566
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
567
- ...
568
-
569
- 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):
570
- """
571
- Specifies the Conda environment for the step.
572
-
573
- Information in this decorator will augment any
574
- attributes set in the `@conda_base` flow-level decorator. Hence,
575
- you can use `@conda_base` to set packages required by all
576
- steps and use `@conda` to specify step-specific overrides.
577
-
578
-
579
- Parameters
580
- ----------
581
- packages : Dict[str, str], default {}
582
- Packages to use for this step. The key is the name of the package
583
- and the value is the version to use.
584
- libraries : Dict[str, str], default {}
585
- Supported for backward compatibility. When used with packages, packages will take precedence.
586
- python : str, optional, default None
587
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
588
- that the version used will correspond to the version of the Python interpreter used to start the run.
589
- disabled : bool, default False
590
- If set to True, disables @conda.
591
- """
592
- ...
593
-
594
- @typing.overload
595
- 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]]]:
706
+ 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):
596
707
  """
597
708
  Specifies a timeout for your step.
598
709
 
@@ -618,139 +729,59 @@ def timeout(*, seconds: int = 0, minutes: int = 0, hours: int = 0) -> typing.Cal
618
729
  ...
619
730
 
620
731
  @typing.overload
621
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
622
- ...
623
-
624
- @typing.overload
625
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
626
- ...
627
-
628
- 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):
732
+ 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]]]:
629
733
  """
630
- Specifies a timeout for your step.
631
-
632
- This decorator is useful if this step may hang indefinitely.
633
-
634
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
635
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
636
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
734
+ Enables checkpointing for a step.
637
735
 
638
- Note that all the values specified in parameters are added together so if you specify
639
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
640
736
 
641
737
 
642
738
  Parameters
643
739
  ----------
644
- seconds : int, default 0
645
- Number of seconds to wait prior to timing out.
646
- minutes : int, default 0
647
- Number of minutes to wait prior to timing out.
648
- hours : int, default 0
649
- Number of hours to wait prior to timing out.
650
- """
651
- ...
652
-
653
- @typing.overload
654
- 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]]]:
655
- """
656
- Specifies the resources needed when executing this step.
657
-
658
- Use `@resources` to specify the resource requirements
659
- independently of the specific compute layer (`@batch`, `@kubernetes`).
660
-
661
- You can choose the compute layer on the command line by executing e.g.
662
- ```
663
- python myflow.py run --with batch
664
- ```
665
- or
666
- ```
667
- python myflow.py run --with kubernetes
668
- ```
669
- which executes the flow on the desired system using the
670
- requirements specified in `@resources`.
671
-
740
+ load_policy : str, default: "fresh"
741
+ The policy for loading the checkpoint. The following policies are supported:
742
+ - "eager": Loads the the latest available checkpoint within the namespace.
743
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
744
+ will be loaded at the start of the task.
745
+ - "none": Do not load any checkpoint
746
+ - "fresh": Loads the lastest checkpoint created within the running Task.
747
+ This mode helps loading checkpoints across various retry attempts of the same task.
748
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
749
+ created within the task will be loaded when the task is retries execution on failure.
672
750
 
673
- Parameters
674
- ----------
675
- cpu : int, default 1
676
- Number of CPUs required for this step.
677
- gpu : int, optional, default None
678
- Number of GPUs required for this step.
679
- disk : int, optional, default None
680
- Disk size (in MB) required for this step. Only applies on Kubernetes.
681
- memory : int, default 4096
682
- Memory size (in MB) required for this step.
683
- shared_memory : int, optional, default None
684
- The value for the size (in MiB) of the /dev/shm volume for this step.
685
- This parameter maps to the `--shm-size` option in Docker.
751
+ temp_dir_root : str, default: None
752
+ The root directory under which `current.checkpoint.directory` will be created.
686
753
  """
687
754
  ...
688
755
 
689
756
  @typing.overload
690
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
757
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
691
758
  ...
692
759
 
693
760
  @typing.overload
694
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
761
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
695
762
  ...
696
763
 
697
- 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):
764
+ 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):
698
765
  """
699
- Specifies the resources needed when executing this step.
700
-
701
- Use `@resources` to specify the resource requirements
702
- independently of the specific compute layer (`@batch`, `@kubernetes`).
703
-
704
- You can choose the compute layer on the command line by executing e.g.
705
- ```
706
- python myflow.py run --with batch
707
- ```
708
- or
709
- ```
710
- python myflow.py run --with kubernetes
711
- ```
712
- which executes the flow on the desired system using the
713
- requirements specified in `@resources`.
714
-
766
+ Enables checkpointing for a step.
715
767
 
716
- Parameters
717
- ----------
718
- cpu : int, default 1
719
- Number of CPUs required for this step.
720
- gpu : int, optional, default None
721
- Number of GPUs required for this step.
722
- disk : int, optional, default None
723
- Disk size (in MB) required for this step. Only applies on Kubernetes.
724
- memory : int, default 4096
725
- Memory size (in MB) required for this step.
726
- shared_memory : int, optional, default None
727
- The value for the size (in MiB) of the /dev/shm volume for this step.
728
- This parameter maps to the `--shm-size` option in Docker.
729
- """
730
- ...
731
-
732
- 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]]]:
733
- """
734
- Decorator that helps cache, version and store models/datasets from huggingface hub.
735
768
 
736
769
 
737
770
  Parameters
738
771
  ----------
739
- temp_dir_root : str, optional
740
- The root directory that will hold the temporary directory where objects will be downloaded.
741
-
742
- load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
743
- The list of repos (models/datasets) to load.
744
-
745
- Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
746
-
747
- - If repo (model/dataset) is not found in the datastore:
748
- - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
749
- - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
750
- - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
772
+ load_policy : str, default: "fresh"
773
+ The policy for loading the checkpoint. The following policies are supported:
774
+ - "eager": Loads the the latest available checkpoint within the namespace.
775
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
776
+ will be loaded at the start of the task.
777
+ - "none": Do not load any checkpoint
778
+ - "fresh": Loads the lastest checkpoint created within the running Task.
779
+ This mode helps loading checkpoints across various retry attempts of the same task.
780
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
781
+ created within the task will be loaded when the task is retries execution on failure.
751
782
 
752
- - If repo is found in the datastore:
753
- - Loads it directly from datastore to local path (can be temporary directory or specified path)
783
+ temp_dir_root : str, default: None
784
+ The root directory under which `current.checkpoint.directory` will be created.
754
785
  """
755
786
  ...
756
787
 
@@ -803,205 +834,138 @@ def card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typ
803
834
  """
804
835
  ...
805
836
 
806
- @typing.overload
807
- def pypi(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
837
+ def ollama(*, models: "list[Ollama]", backend: 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]]]:
808
838
  """
809
- Specifies the PyPI packages for the step.
839
+ This decorator is used to run Ollama APIs as Metaflow task sidecars.
810
840
 
811
- Information in this decorator will augment any
812
- attributes set in the `@pyi_base` flow-level decorator. Hence,
813
- you can use `@pypi_base` to set packages required by all
814
- steps and use `@pypi` to specify step-specific overrides.
841
+ User code call
842
+ -----------
843
+ @ollama(
844
+ models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
845
+ backend='local'
846
+ )
847
+
848
+ Valid backend options
849
+ ---------------------
850
+ - 'local': Run as a separate process on the local task machine.
851
+ - (TODO) 'managed': Outerbounds hosts and selects compute provider.
852
+ - (TODO) 'remote': Spin up separate instance to serve Ollama models.
853
+
854
+ Valid model options
855
+ ----------------
856
+ - 'llama3.2'
857
+ - 'llama3.3'
858
+ - any model here https://ollama.com/search
815
859
 
816
860
 
817
861
  Parameters
818
862
  ----------
819
- packages : Dict[str, str], default: {}
820
- Packages to use for this step. The key is the name of the package
821
- and the value is the version to use.
822
- python : str, optional, default: None
823
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
824
- that the version used will correspond to the version of the Python interpreter used to start the run.
863
+ models: list[Ollama]
864
+ List of Ollama containers running models in sidecars.
865
+ backend: str
866
+ Determines where and how to run the Ollama process.
825
867
  """
826
868
  ...
827
869
 
828
870
  @typing.overload
829
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
871
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
872
+ """
873
+ Decorator prototype for all step decorators. This function gets specialized
874
+ and imported for all decorators types by _import_plugin_decorators().
875
+ """
830
876
  ...
831
877
 
832
878
  @typing.overload
833
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
879
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
834
880
  ...
835
881
 
836
- def pypi(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
882
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
837
883
  """
838
- Specifies the PyPI packages for the step.
839
-
840
- Information in this decorator will augment any
841
- attributes set in the `@pyi_base` flow-level decorator. Hence,
842
- you can use `@pypi_base` to set packages required by all
843
- steps and use `@pypi` to specify step-specific overrides.
844
-
845
-
846
- Parameters
847
- ----------
848
- packages : Dict[str, str], default: {}
849
- Packages to use for this step. The key is the name of the package
850
- and the value is the version to use.
851
- python : str, optional, default: None
852
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
853
- that the version used will correspond to the version of the Python interpreter used to start the run.
854
- """
855
- ...
856
-
857
- def nim(*, models: "list[NIM]", backend: 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]]]:
858
- """
859
- This decorator is used to run NIM containers in Metaflow tasks as sidecars.
860
-
861
- User code call
862
- -----------
863
- @nim(
864
- models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
865
- backend='managed'
866
- )
867
-
868
- Valid backend options
869
- ---------------------
870
- - 'managed': Outerbounds selects a compute provider based on the model.
871
-
872
- Valid model options
873
- ----------------
874
- - 'meta/llama3-8b-instruct': 8B parameter model
875
- - 'meta/llama3-70b-instruct': 70B parameter model
876
- - any model here: https://nvcf.ngc.nvidia.com/functions?filter=nvidia-functions
877
-
878
-
879
- Parameters
880
- ----------
881
- models: list[NIM]
882
- List of NIM containers running models in sidecars.
883
- backend: str
884
- Compute provider to run the NIM container.
885
- queue_timeout : int
886
- Time to keep the job in NVCF's queue.
884
+ Decorator prototype for all step decorators. This function gets specialized
885
+ and imported for all decorators types by _import_plugin_decorators().
887
886
  """
888
887
  ...
889
888
 
890
889
  @typing.overload
891
- 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]]]:
890
+ 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]]]:
892
891
  """
893
- Specifies that the step will success under all circumstances.
894
-
895
- The decorator will create an optional artifact, specified by `var`, which
896
- contains the exception raised. You can use it to detect the presence
897
- of errors, indicating that all happy-path artifacts produced by the step
898
- are missing.
892
+ Specifies secrets to be retrieved and injected as environment variables prior to
893
+ the execution of a step.
899
894
 
900
895
 
901
896
  Parameters
902
897
  ----------
903
- var : str, optional, default None
904
- Name of the artifact in which to store the caught exception.
905
- If not specified, the exception is not stored.
906
- print_exception : bool, default True
907
- Determines whether or not the exception is printed to
908
- stdout when caught.
898
+ sources : List[Union[str, Dict[str, Any]]], default: []
899
+ List of secret specs, defining how the secrets are to be retrieved
909
900
  """
910
901
  ...
911
902
 
912
903
  @typing.overload
913
- def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
904
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
914
905
  ...
915
906
 
916
907
  @typing.overload
917
- def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
908
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
918
909
  ...
919
910
 
920
- 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):
911
+ 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]]] = []):
921
912
  """
922
- Specifies that the step will success under all circumstances.
923
-
924
- The decorator will create an optional artifact, specified by `var`, which
925
- contains the exception raised. You can use it to detect the presence
926
- of errors, indicating that all happy-path artifacts produced by the step
927
- are missing.
913
+ Specifies secrets to be retrieved and injected as environment variables prior to
914
+ the execution of a step.
928
915
 
929
916
 
930
917
  Parameters
931
918
  ----------
932
- var : str, optional, default None
933
- Name of the artifact in which to store the caught exception.
934
- If not specified, the exception is not stored.
935
- print_exception : bool, default True
936
- Determines whether or not the exception is printed to
937
- stdout when caught.
919
+ sources : List[Union[str, Dict[str, Any]]], default: []
920
+ List of secret specs, defining how the secrets are to be retrieved
938
921
  """
939
922
  ...
940
923
 
941
- 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]]]:
924
+ @typing.overload
925
+ 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]]]:
942
926
  """
943
- Specifies that this step should execute on DGX cloud.
927
+ Specifies environment variables to be set prior to the execution of a step.
944
928
 
945
929
 
946
930
  Parameters
947
931
  ----------
948
- gpu : int
949
- Number of GPUs to use.
950
- gpu_type : str
951
- Type of Nvidia GPU to use.
952
- queue_timeout : int
953
- Time to keep the job in NVCF's queue.
932
+ vars : Dict[str, str], default {}
933
+ Dictionary of environment variables to set.
954
934
  """
955
935
  ...
956
936
 
957
937
  @typing.overload
958
- 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]]:
938
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
939
+ ...
940
+
941
+ @typing.overload
942
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
943
+ ...
944
+
945
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
959
946
  """
960
- Specifies the times when the flow should be run when running on a
961
- production scheduler.
947
+ Specifies environment variables to be set prior to the execution of a step.
962
948
 
963
949
 
964
950
  Parameters
965
951
  ----------
966
- hourly : bool, default False
967
- Run the workflow hourly.
968
- daily : bool, default True
969
- Run the workflow daily.
970
- weekly : bool, default False
971
- Run the workflow weekly.
972
- cron : str, optional, default None
973
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
974
- specified by this expression.
975
- timezone : str, optional, default None
976
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
977
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
952
+ vars : Dict[str, str], default {}
953
+ Dictionary of environment variables to set.
978
954
  """
979
955
  ...
980
956
 
981
- @typing.overload
982
- def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
983
- ...
984
-
985
- 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):
957
+ 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]]]:
986
958
  """
987
- Specifies the times when the flow should be run when running on a
988
- production scheduler.
959
+ Specifies that this step is used to deploy an instance of the app.
960
+ Requires that self.app_name, self.app_port, self.entrypoint and self.deployDir is set.
989
961
 
990
962
 
991
963
  Parameters
992
964
  ----------
993
- hourly : bool, default False
994
- Run the workflow hourly.
995
- daily : bool, default True
996
- Run the workflow daily.
997
- weekly : bool, default False
998
- Run the workflow weekly.
999
- cron : str, optional, default None
1000
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1001
- specified by this expression.
1002
- timezone : str, optional, default None
1003
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1004
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
965
+ app_port : int
966
+ Number of GPUs to use.
967
+ app_name : str
968
+ Name of the app to deploy.
1005
969
  """
1006
970
  ...
1007
971
 
@@ -1019,211 +983,24 @@ def project(*, name: str, branch: typing.Optional[str] = None, production: bool
1019
983
  Project name. Make sure that the name is unique amongst all
1020
984
  projects that use the same production scheduler. The name may
1021
985
  contain only lowercase alphanumeric characters and underscores.
1022
-
1023
- branch : Optional[str], default None
1024
- The branch to use. If not specified, the branch is set to
1025
- `user.<username>` unless `production` is set to `True`. This can
1026
- also be set on the command line using `--branch` as a top-level option.
1027
- It is an error to specify `branch` in the decorator and on the command line.
1028
-
1029
- production : bool, default False
1030
- Whether or not the branch is the production branch. This can also be set on the
1031
- command line using `--production` as a top-level option. It is an error to specify
1032
- `production` in the decorator and on the command line.
1033
- The project branch name will be:
1034
- - if `branch` is specified:
1035
- - if `production` is True: `prod.<branch>`
1036
- - if `production` is False: `test.<branch>`
1037
- - if `branch` is not specified:
1038
- - if `production` is True: `prod`
1039
- - if `production` is False: `user.<username>`
1040
- """
1041
- ...
1042
-
1043
- 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]]:
1044
- """
1045
- 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.
1046
- This decorator only works when a flow is scheduled on Airflow and is compiled using `airflow create`. More than one `@airflow_external_task_sensor` can be added as a flow decorators. Adding more than one decorator will ensure that `start` step starts only after all sensors finish.
1047
-
1048
-
1049
- Parameters
1050
- ----------
1051
- timeout : int
1052
- Time, in seconds before the task times out and fails. (Default: 3600)
1053
- poke_interval : int
1054
- Time in seconds that the job should wait in between each try. (Default: 60)
1055
- mode : str
1056
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1057
- exponential_backoff : bool
1058
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1059
- pool : str
1060
- the slot pool this task should run in,
1061
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1062
- soft_fail : bool
1063
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1064
- name : str
1065
- Name of the sensor on Airflow
1066
- description : str
1067
- Description of sensor in the Airflow UI
1068
- external_dag_id : str
1069
- The dag_id that contains the task you want to wait for.
1070
- external_task_ids : List[str]
1071
- The list of task_ids that you want to wait for.
1072
- If None (default value) the sensor waits for the DAG. (Default: None)
1073
- allowed_states : List[str]
1074
- Iterable of allowed states, (Default: ['success'])
1075
- failed_states : List[str]
1076
- Iterable of failed or dis-allowed states. (Default: None)
1077
- execution_delta : datetime.timedelta
1078
- time difference with the previous execution to look at,
1079
- the default is the same logical date as the current task or DAG. (Default: None)
1080
- check_existence: bool
1081
- Set to True to check if the external task exists or check if
1082
- the DAG to wait for exists. (Default: True)
1083
- """
1084
- ...
1085
-
1086
- @typing.overload
1087
- def trigger(*, event: typing.Union[str, typing.Dict[str, typing.Any], None] = None, events: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], options: typing.Dict[str, typing.Any] = {}) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1088
- """
1089
- Specifies the event(s) that this flow depends on.
1090
-
1091
- ```
1092
- @trigger(event='foo')
1093
- ```
1094
- or
1095
- ```
1096
- @trigger(events=['foo', 'bar'])
1097
- ```
1098
-
1099
- Additionally, you can specify the parameter mappings
1100
- to map event payload to Metaflow parameters for the flow.
1101
- ```
1102
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1103
- ```
1104
- or
1105
- ```
1106
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1107
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1108
- ```
1109
-
1110
- 'parameters' can also be a list of strings and tuples like so:
1111
- ```
1112
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1113
- ```
1114
- This is equivalent to:
1115
- ```
1116
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1117
- ```
1118
-
1119
-
1120
- Parameters
1121
- ----------
1122
- event : Union[str, Dict[str, Any]], optional, default None
1123
- Event dependency for this flow.
1124
- events : List[Union[str, Dict[str, Any]]], default []
1125
- Events dependency for this flow.
1126
- options : Dict[str, Any], default {}
1127
- Backend-specific configuration for tuning eventing behavior.
1128
- """
1129
- ...
1130
-
1131
- @typing.overload
1132
- def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1133
- ...
1134
-
1135
- def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: typing.Union[str, typing.Dict[str, typing.Any], None] = None, events: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], options: typing.Dict[str, typing.Any] = {}):
1136
- """
1137
- Specifies the event(s) that this flow depends on.
1138
-
1139
- ```
1140
- @trigger(event='foo')
1141
- ```
1142
- or
1143
- ```
1144
- @trigger(events=['foo', 'bar'])
1145
- ```
1146
-
1147
- Additionally, you can specify the parameter mappings
1148
- to map event payload to Metaflow parameters for the flow.
1149
- ```
1150
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1151
- ```
1152
- or
1153
- ```
1154
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1155
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1156
- ```
1157
-
1158
- 'parameters' can also be a list of strings and tuples like so:
1159
- ```
1160
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1161
- ```
1162
- This is equivalent to:
1163
- ```
1164
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1165
- ```
1166
-
1167
-
1168
- Parameters
1169
- ----------
1170
- event : Union[str, Dict[str, Any]], optional, default None
1171
- Event dependency for this flow.
1172
- events : List[Union[str, Dict[str, Any]]], default []
1173
- Events dependency for this flow.
1174
- options : Dict[str, Any], default {}
1175
- Backend-specific configuration for tuning eventing behavior.
1176
- """
1177
- ...
1178
-
1179
- @typing.overload
1180
- 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]]:
1181
- """
1182
- Specifies the Conda environment for all steps of the flow.
1183
-
1184
- Use `@conda_base` to set common libraries required by all
1185
- steps and use `@conda` to specify step-specific additions.
1186
-
1187
-
1188
- Parameters
1189
- ----------
1190
- packages : Dict[str, str], default {}
1191
- Packages to use for this flow. The key is the name of the package
1192
- and the value is the version to use.
1193
- libraries : Dict[str, str], default {}
1194
- Supported for backward compatibility. When used with packages, packages will take precedence.
1195
- python : str, optional, default None
1196
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1197
- that the version used will correspond to the version of the Python interpreter used to start the run.
1198
- disabled : bool, default False
1199
- If set to True, disables Conda.
1200
- """
1201
- ...
1202
-
1203
- @typing.overload
1204
- def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1205
- ...
1206
-
1207
- 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):
1208
- """
1209
- Specifies the Conda environment for all steps of the flow.
1210
-
1211
- Use `@conda_base` to set common libraries required by all
1212
- steps and use `@conda` to specify step-specific additions.
1213
-
1214
-
1215
- Parameters
1216
- ----------
1217
- packages : Dict[str, str], default {}
1218
- Packages to use for this flow. The key is the name of the package
1219
- and the value is the version to use.
1220
- libraries : Dict[str, str], default {}
1221
- Supported for backward compatibility. When used with packages, packages will take precedence.
1222
- python : str, optional, default None
1223
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1224
- that the version used will correspond to the version of the Python interpreter used to start the run.
1225
- disabled : bool, default False
1226
- If set to True, disables Conda.
986
+
987
+ branch : Optional[str], default None
988
+ The branch to use. If not specified, the branch is set to
989
+ `user.<username>` unless `production` is set to `True`. This can
990
+ also be set on the command line using `--branch` as a top-level option.
991
+ It is an error to specify `branch` in the decorator and on the command line.
992
+
993
+ production : bool, default False
994
+ Whether or not the branch is the production branch. This can also be set on the
995
+ command line using `--production` as a top-level option. It is an error to specify
996
+ `production` in the decorator and on the command line.
997
+ The project branch name will be:
998
+ - if `branch` is specified:
999
+ - if `production` is True: `prod.<branch>`
1000
+ - if `production` is False: `test.<branch>`
1001
+ - if `branch` is not specified:
1002
+ - if `production` is True: `prod`
1003
+ - if `production` is False: `user.<username>`
1227
1004
  """
1228
1005
  ...
1229
1006
 
@@ -1371,6 +1148,100 @@ def airflow_s3_key_sensor(*, timeout: int, poke_interval: int, mode: str, expone
1371
1148
  """
1372
1149
  ...
1373
1150
 
1151
+ @typing.overload
1152
+ 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]]:
1153
+ """
1154
+ Specifies the Conda environment for all steps of the flow.
1155
+
1156
+ Use `@conda_base` to set common libraries required by all
1157
+ steps and use `@conda` to specify step-specific additions.
1158
+
1159
+
1160
+ Parameters
1161
+ ----------
1162
+ packages : Dict[str, str], default {}
1163
+ Packages to use for this flow. The key is the name of the package
1164
+ and the value is the version to use.
1165
+ libraries : Dict[str, str], default {}
1166
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1167
+ python : str, optional, default None
1168
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1169
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1170
+ disabled : bool, default False
1171
+ If set to True, disables Conda.
1172
+ """
1173
+ ...
1174
+
1175
+ @typing.overload
1176
+ def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1177
+ ...
1178
+
1179
+ 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):
1180
+ """
1181
+ Specifies the Conda environment for all steps of the flow.
1182
+
1183
+ Use `@conda_base` to set common libraries required by all
1184
+ steps and use `@conda` to specify step-specific additions.
1185
+
1186
+
1187
+ Parameters
1188
+ ----------
1189
+ packages : Dict[str, str], default {}
1190
+ Packages to use for this flow. The key is the name of the package
1191
+ and the value is the version to use.
1192
+ libraries : Dict[str, str], default {}
1193
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1194
+ python : str, optional, default None
1195
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1196
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1197
+ disabled : bool, default False
1198
+ If set to True, disables Conda.
1199
+ """
1200
+ ...
1201
+
1202
+ 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]]:
1203
+ """
1204
+ 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.
1205
+ This decorator only works when a flow is scheduled on Airflow and is compiled using `airflow create`. More than one `@airflow_external_task_sensor` can be added as a flow decorators. Adding more than one decorator will ensure that `start` step starts only after all sensors finish.
1206
+
1207
+
1208
+ Parameters
1209
+ ----------
1210
+ timeout : int
1211
+ Time, in seconds before the task times out and fails. (Default: 3600)
1212
+ poke_interval : int
1213
+ Time in seconds that the job should wait in between each try. (Default: 60)
1214
+ mode : str
1215
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1216
+ exponential_backoff : bool
1217
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1218
+ pool : str
1219
+ the slot pool this task should run in,
1220
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1221
+ soft_fail : bool
1222
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1223
+ name : str
1224
+ Name of the sensor on Airflow
1225
+ description : str
1226
+ Description of sensor in the Airflow UI
1227
+ external_dag_id : str
1228
+ The dag_id that contains the task you want to wait for.
1229
+ external_task_ids : List[str]
1230
+ The list of task_ids that you want to wait for.
1231
+ If None (default value) the sensor waits for the DAG. (Default: None)
1232
+ allowed_states : List[str]
1233
+ Iterable of allowed states, (Default: ['success'])
1234
+ failed_states : List[str]
1235
+ Iterable of failed or dis-allowed states. (Default: None)
1236
+ execution_delta : datetime.timedelta
1237
+ time difference with the previous execution to look at,
1238
+ the default is the same logical date as the current task or DAG. (Default: None)
1239
+ check_existence: bool
1240
+ Set to True to check if the external task exists or check if
1241
+ the DAG to wait for exists. (Default: True)
1242
+ """
1243
+ ...
1244
+
1374
1245
  @typing.overload
1375
1246
  def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1376
1247
  """
@@ -1412,5 +1283,149 @@ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packag
1412
1283
  """
1413
1284
  ...
1414
1285
 
1286
+ @typing.overload
1287
+ def trigger(*, event: typing.Union[str, typing.Dict[str, typing.Any], None] = None, events: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], options: typing.Dict[str, typing.Any] = {}) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1288
+ """
1289
+ Specifies the event(s) that this flow depends on.
1290
+
1291
+ ```
1292
+ @trigger(event='foo')
1293
+ ```
1294
+ or
1295
+ ```
1296
+ @trigger(events=['foo', 'bar'])
1297
+ ```
1298
+
1299
+ Additionally, you can specify the parameter mappings
1300
+ to map event payload to Metaflow parameters for the flow.
1301
+ ```
1302
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1303
+ ```
1304
+ or
1305
+ ```
1306
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1307
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1308
+ ```
1309
+
1310
+ 'parameters' can also be a list of strings and tuples like so:
1311
+ ```
1312
+ @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1313
+ ```
1314
+ This is equivalent to:
1315
+ ```
1316
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1317
+ ```
1318
+
1319
+
1320
+ Parameters
1321
+ ----------
1322
+ event : Union[str, Dict[str, Any]], optional, default None
1323
+ Event dependency for this flow.
1324
+ events : List[Union[str, Dict[str, Any]]], default []
1325
+ Events dependency for this flow.
1326
+ options : Dict[str, Any], default {}
1327
+ Backend-specific configuration for tuning eventing behavior.
1328
+ """
1329
+ ...
1330
+
1331
+ @typing.overload
1332
+ def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1333
+ ...
1334
+
1335
+ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: typing.Union[str, typing.Dict[str, typing.Any], None] = None, events: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], options: typing.Dict[str, typing.Any] = {}):
1336
+ """
1337
+ Specifies the event(s) that this flow depends on.
1338
+
1339
+ ```
1340
+ @trigger(event='foo')
1341
+ ```
1342
+ or
1343
+ ```
1344
+ @trigger(events=['foo', 'bar'])
1345
+ ```
1346
+
1347
+ Additionally, you can specify the parameter mappings
1348
+ to map event payload to Metaflow parameters for the flow.
1349
+ ```
1350
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1351
+ ```
1352
+ or
1353
+ ```
1354
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1355
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1356
+ ```
1357
+
1358
+ 'parameters' can also be a list of strings and tuples like so:
1359
+ ```
1360
+ @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1361
+ ```
1362
+ This is equivalent to:
1363
+ ```
1364
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1365
+ ```
1366
+
1367
+
1368
+ Parameters
1369
+ ----------
1370
+ event : Union[str, Dict[str, Any]], optional, default None
1371
+ Event dependency for this flow.
1372
+ events : List[Union[str, Dict[str, Any]]], default []
1373
+ Events dependency for this flow.
1374
+ options : Dict[str, Any], default {}
1375
+ Backend-specific configuration for tuning eventing behavior.
1376
+ """
1377
+ ...
1378
+
1379
+ @typing.overload
1380
+ 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]]:
1381
+ """
1382
+ Specifies the times when the flow should be run when running on a
1383
+ production scheduler.
1384
+
1385
+
1386
+ Parameters
1387
+ ----------
1388
+ hourly : bool, default False
1389
+ Run the workflow hourly.
1390
+ daily : bool, default True
1391
+ Run the workflow daily.
1392
+ weekly : bool, default False
1393
+ Run the workflow weekly.
1394
+ cron : str, optional, default None
1395
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1396
+ specified by this expression.
1397
+ timezone : str, optional, default None
1398
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1399
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1400
+ """
1401
+ ...
1402
+
1403
+ @typing.overload
1404
+ def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1405
+ ...
1406
+
1407
+ 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):
1408
+ """
1409
+ Specifies the times when the flow should be run when running on a
1410
+ production scheduler.
1411
+
1412
+
1413
+ Parameters
1414
+ ----------
1415
+ hourly : bool, default False
1416
+ Run the workflow hourly.
1417
+ daily : bool, default True
1418
+ Run the workflow daily.
1419
+ weekly : bool, default False
1420
+ Run the workflow weekly.
1421
+ cron : str, optional, default None
1422
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1423
+ specified by this expression.
1424
+ timezone : str, optional, default None
1425
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1426
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1427
+ """
1428
+ ...
1429
+
1415
1430
  pkg_name: str
1416
1431