ob-metaflow-stubs 6.0.3.142__py2.py3-none-any.whl → 6.0.3.143rc0__py2.py3-none-any.whl

Sign up to get free protection for your applications and to get access to all the features.
Files changed (194) hide show
  1. metaflow-stubs/__init__.pyi +547 -516
  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 +3 -3
  9. metaflow-stubs/events.pyi +2 -2
  10. metaflow-stubs/exception.pyi +2 -2
  11. metaflow-stubs/flowspec.pyi +6 -6
  12. metaflow-stubs/generated_for.txt +1 -1
  13. metaflow-stubs/includefile.pyi +4 -4
  14. metaflow-stubs/info_file.pyi +2 -2
  15. metaflow-stubs/metadata_provider/__init__.pyi +2 -2
  16. metaflow-stubs/metadata_provider/heartbeat.pyi +2 -2
  17. metaflow-stubs/metadata_provider/metadata.pyi +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 +80 -80
  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 +2 -2
  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 +3 -3
  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 +5 -5
  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 +7 -5
  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 +7 -4
  36. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/decorator.pyi +6 -5
  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 +4 -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/context.pyi +79 -0
  42. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/core.pyi +4 -8
  43. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/decorator.pyi +46 -0
  44. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/exceptions.pyi +2 -2
  45. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/task_utils.pyi +11 -4
  46. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/utils.pyi +2 -2
  47. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastructures.pyi +3 -3
  48. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/exceptions.pyi +2 -2
  49. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/__init__.pyi +2 -2
  50. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/decorator.pyi +2 -2
  51. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/__init__.pyi +2 -2
  52. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/core.pyi +2 -2
  53. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/exceptions.pyi +2 -2
  54. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/model_storage.pyi +4 -4
  55. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/__init__.pyi +2 -2
  56. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/flowspec_utils.pyi +2 -2
  57. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/general.pyi +2 -2
  58. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/identity_utils.pyi +2 -2
  59. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/__init__.pyi +2 -2
  60. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/base.pyi +2 -2
  61. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/tar.pyi +3 -3
  62. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/tar_utils.pyi +3 -3
  63. metaflow-stubs/mf_extensions/outerbounds/__init__.pyi +2 -2
  64. metaflow-stubs/mf_extensions/outerbounds/plugins/__init__.pyi +2 -2
  65. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/__init__.pyi +2 -2
  66. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/snowflake.pyi +2 -2
  67. metaflow-stubs/mf_extensions/outerbounds/profilers/__init__.pyi +2 -2
  68. metaflow-stubs/mf_extensions/outerbounds/profilers/gpu.pyi +2 -2
  69. metaflow-stubs/mf_extensions/outerbounds/remote_config.pyi +2 -2
  70. metaflow-stubs/mf_extensions/outerbounds/toplevel/__init__.pyi +2 -2
  71. metaflow-stubs/mf_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.pyi +2 -2
  72. metaflow-stubs/multicore_utils.pyi +2 -2
  73. metaflow-stubs/parameters.pyi +4 -4
  74. metaflow-stubs/plugins/__init__.pyi +13 -13
  75. metaflow-stubs/plugins/airflow/__init__.pyi +2 -2
  76. metaflow-stubs/plugins/airflow/airflow_utils.pyi +2 -2
  77. metaflow-stubs/plugins/airflow/exception.pyi +2 -2
  78. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +2 -2
  79. metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi +2 -2
  80. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +2 -2
  81. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +2 -2
  82. metaflow-stubs/plugins/argo/__init__.pyi +2 -2
  83. metaflow-stubs/plugins/argo/argo_client.pyi +2 -2
  84. metaflow-stubs/plugins/argo/argo_events.pyi +2 -2
  85. metaflow-stubs/plugins/argo/argo_workflows.pyi +2 -2
  86. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +3 -3
  87. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +4 -4
  88. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +4 -4
  89. metaflow-stubs/plugins/aws/__init__.pyi +2 -2
  90. metaflow-stubs/plugins/aws/aws_client.pyi +2 -2
  91. metaflow-stubs/plugins/aws/aws_utils.pyi +2 -2
  92. metaflow-stubs/plugins/aws/batch/__init__.pyi +2 -2
  93. metaflow-stubs/plugins/aws/batch/batch.pyi +2 -2
  94. metaflow-stubs/plugins/aws/batch/batch_client.pyi +2 -2
  95. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +2 -2
  96. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +2 -2
  97. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +4 -4
  98. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +2 -2
  99. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +2 -2
  100. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +2 -2
  101. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +2 -2
  102. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +2 -2
  103. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer.pyi +4 -4
  104. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer_objects.pyi +4 -4
  105. metaflow-stubs/plugins/azure/__init__.pyi +2 -2
  106. metaflow-stubs/plugins/azure/azure_credential.pyi +2 -2
  107. metaflow-stubs/plugins/azure/azure_exceptions.pyi +2 -2
  108. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +4 -4
  109. metaflow-stubs/plugins/azure/azure_utils.pyi +2 -2
  110. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +2 -2
  111. metaflow-stubs/plugins/azure/includefile_support.pyi +2 -2
  112. metaflow-stubs/plugins/cards/__init__.pyi +2 -2
  113. metaflow-stubs/plugins/cards/card_client.pyi +3 -3
  114. metaflow-stubs/plugins/cards/card_creator.pyi +2 -2
  115. metaflow-stubs/plugins/cards/card_datastore.pyi +2 -2
  116. metaflow-stubs/plugins/cards/card_decorator.pyi +2 -2
  117. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +2 -2
  118. metaflow-stubs/plugins/cards/card_modules/basic.pyi +2 -2
  119. metaflow-stubs/plugins/cards/card_modules/card.pyi +2 -2
  120. metaflow-stubs/plugins/cards/card_modules/components.pyi +4 -4
  121. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +2 -2
  122. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +2 -2
  123. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +2 -2
  124. metaflow-stubs/plugins/cards/card_resolver.pyi +2 -2
  125. metaflow-stubs/plugins/cards/component_serializer.pyi +2 -2
  126. metaflow-stubs/plugins/cards/exception.pyi +2 -2
  127. metaflow-stubs/plugins/catch_decorator.pyi +2 -2
  128. metaflow-stubs/plugins/datatools/__init__.pyi +2 -2
  129. metaflow-stubs/plugins/datatools/local.pyi +2 -2
  130. metaflow-stubs/plugins/datatools/s3/__init__.pyi +2 -2
  131. metaflow-stubs/plugins/datatools/s3/s3.pyi +4 -4
  132. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +2 -2
  133. metaflow-stubs/plugins/datatools/s3/s3util.pyi +2 -2
  134. metaflow-stubs/plugins/debug_logger.pyi +2 -2
  135. metaflow-stubs/plugins/debug_monitor.pyi +2 -2
  136. metaflow-stubs/plugins/environment_decorator.pyi +2 -2
  137. metaflow-stubs/plugins/events_decorator.pyi +2 -2
  138. metaflow-stubs/plugins/frameworks/__init__.pyi +2 -2
  139. metaflow-stubs/plugins/frameworks/pytorch.pyi +2 -2
  140. metaflow-stubs/plugins/gcp/__init__.pyi +2 -2
  141. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +4 -4
  142. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +2 -2
  143. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +2 -2
  144. metaflow-stubs/plugins/gcp/gs_utils.pyi +2 -2
  145. metaflow-stubs/plugins/gcp/includefile_support.pyi +2 -2
  146. metaflow-stubs/plugins/kubernetes/__init__.pyi +2 -2
  147. metaflow-stubs/plugins/kubernetes/kube_utils.pyi +3 -3
  148. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +4 -4
  149. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +2 -2
  150. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +5 -2
  151. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +2 -2
  152. metaflow-stubs/plugins/kubernetes/spot_monitor_sidecar.pyi +2 -2
  153. metaflow-stubs/plugins/parallel_decorator.pyi +2 -2
  154. metaflow-stubs/plugins/perimeters.pyi +2 -2
  155. metaflow-stubs/plugins/project_decorator.pyi +2 -2
  156. metaflow-stubs/plugins/pypi/__init__.pyi +3 -3
  157. metaflow-stubs/plugins/pypi/conda_decorator.pyi +2 -2
  158. metaflow-stubs/plugins/pypi/conda_environment.pyi +5 -5
  159. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +2 -2
  160. metaflow-stubs/plugins/pypi/pypi_environment.pyi +2 -2
  161. metaflow-stubs/plugins/pypi/utils.pyi +2 -2
  162. metaflow-stubs/plugins/resources_decorator.pyi +2 -2
  163. metaflow-stubs/plugins/retry_decorator.pyi +2 -2
  164. metaflow-stubs/plugins/secrets/__init__.pyi +3 -3
  165. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +4 -4
  166. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +2 -2
  167. metaflow-stubs/plugins/snowflake/__init__.pyi +2 -2
  168. metaflow-stubs/plugins/storage_executor.pyi +2 -2
  169. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +3 -3
  170. metaflow-stubs/plugins/timeout_decorator.pyi +2 -2
  171. metaflow-stubs/profilers/__init__.pyi +2 -2
  172. metaflow-stubs/pylint_wrapper.pyi +2 -2
  173. metaflow-stubs/runner/__init__.pyi +2 -2
  174. metaflow-stubs/runner/deployer.pyi +5 -5
  175. metaflow-stubs/runner/deployer_impl.pyi +3 -3
  176. metaflow-stubs/runner/metaflow_runner.pyi +3 -3
  177. metaflow-stubs/runner/nbdeploy.pyi +2 -2
  178. metaflow-stubs/runner/nbrun.pyi +2 -2
  179. metaflow-stubs/runner/subprocess_manager.pyi +2 -2
  180. metaflow-stubs/runner/utils.pyi +3 -3
  181. metaflow-stubs/system/__init__.pyi +2 -2
  182. metaflow-stubs/system/system_logger.pyi +2 -2
  183. metaflow-stubs/system/system_monitor.pyi +2 -2
  184. metaflow-stubs/tagging_util.pyi +2 -2
  185. metaflow-stubs/tuple_util.pyi +2 -2
  186. metaflow-stubs/user_configs/__init__.pyi +2 -2
  187. metaflow-stubs/user_configs/config_decorators.pyi +6 -6
  188. metaflow-stubs/user_configs/config_options.pyi +3 -3
  189. metaflow-stubs/user_configs/config_parameters.pyi +7 -7
  190. {ob_metaflow_stubs-6.0.3.142.dist-info → ob_metaflow_stubs-6.0.3.143rc0.dist-info}/METADATA +1 -1
  191. ob_metaflow_stubs-6.0.3.143rc0.dist-info/RECORD +194 -0
  192. ob_metaflow_stubs-6.0.3.142.dist-info/RECORD +0 -192
  193. {ob_metaflow_stubs-6.0.3.142.dist-info → ob_metaflow_stubs-6.0.3.143rc0.dist-info}/WHEEL +0 -0
  194. {ob_metaflow_stubs-6.0.3.142.dist-info → ob_metaflow_stubs-6.0.3.143rc0.dist-info}/top_level.txt +0 -0
@@ -1,15 +1,15 @@
1
1
  ######################################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
- # MF version: 2.14.0.1+obcheckpoint(0.1.8);ob(v1) #
4
- # Generated on 2025-02-13T17:42:31.232300 #
3
+ # MF version: 2.14.0.2rc0+obcheckpoint(0.1.8.1rc0);ob(v1) #
4
+ # Generated on 2025-02-17T19:27:36.178389 #
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
 
@@ -35,8 +35,8 @@ from .user_configs.config_parameters import ConfigValue as ConfigValue
35
35
  from .user_configs.config_parameters import config_expr as config_expr
36
36
  from .user_configs.config_decorators import CustomFlowDecorator as CustomFlowDecorator
37
37
  from .user_configs.config_decorators import CustomStepDecorator as CustomStepDecorator
38
- from . import cards as cards
39
38
  from . import tuple_util as tuple_util
39
+ from . import cards as cards
40
40
  from . import events as events
41
41
  from . import runner as runner
42
42
  from . import plugins as plugins
@@ -63,6 +63,7 @@ from .runner.deployer import DeployedFlow as DeployedFlow
63
63
  from .runner.nbdeploy import NBDeployer as NBDeployer
64
64
  from .mf_extensions.obcheckpoint.plugins.machine_learning_utilities.checkpoints.final_api import Checkpoint as Checkpoint
65
65
  from .mf_extensions.obcheckpoint.plugins.machine_learning_utilities.datastructures import load_model as load_model
66
+ from .mf_extensions.obcheckpoint.plugins.machine_learning_utilities.datastore.context import using_artifact_store as using_artifact_store
66
67
  from .mf_extensions.outerbounds.toplevel.global_aliases_for_metaflow_package import get_aws_client as get_aws_client
67
68
  from .mf_extensions.outerbounds.plugins.snowflake.snowflake import Snowflake as Snowflake
68
69
  from . import cli_components as cli_components
@@ -146,279 +147,247 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
146
147
  """
147
148
  ...
148
149
 
149
- @typing.overload
150
- 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]]]:
150
+ 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]]]:
151
151
  """
152
- Specifies the resources needed when executing this step.
153
-
154
- Use `@resources` to specify the resource requirements
155
- independently of the specific compute layer (`@batch`, `@kubernetes`).
156
-
157
- You can choose the compute layer on the command line by executing e.g.
158
- ```
159
- python myflow.py run --with batch
160
- ```
161
- or
162
- ```
163
- python myflow.py run --with kubernetes
164
- ```
165
- which executes the flow on the desired system using the
166
- requirements specified in `@resources`.
152
+ Decorator that helps cache, version and store models/datasets from huggingface hub.
167
153
 
168
154
 
169
155
  Parameters
170
156
  ----------
171
- cpu : int, default 1
172
- Number of CPUs required for this step.
173
- gpu : int, optional, default None
174
- Number of GPUs required for this step.
175
- disk : int, optional, default None
176
- Disk size (in MB) required for this step. Only applies on Kubernetes.
177
- memory : int, default 4096
178
- Memory size (in MB) required for this step.
179
- shared_memory : int, optional, default None
180
- The value for the size (in MiB) of the /dev/shm volume for this step.
181
- This parameter maps to the `--shm-size` option in Docker.
182
- """
183
- ...
184
-
185
- @typing.overload
186
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
187
- ...
188
-
189
- @typing.overload
190
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
191
- ...
192
-
193
- 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):
194
- """
195
- Specifies the resources needed when executing this step.
157
+ temp_dir_root : str, optional
158
+ The root directory that will hold the temporary directory where objects will be downloaded.
196
159
 
197
- Use `@resources` to specify the resource requirements
198
- independently of the specific compute layer (`@batch`, `@kubernetes`).
160
+ load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
161
+ The list of repos (models/datasets) to load.
199
162
 
200
- You can choose the compute layer on the command line by executing e.g.
201
- ```
202
- python myflow.py run --with batch
203
- ```
204
- or
205
- ```
206
- python myflow.py run --with kubernetes
207
- ```
208
- which executes the flow on the desired system using the
209
- requirements specified in `@resources`.
163
+ Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
210
164
 
165
+ - If repo (model/dataset) is not found in the datastore:
166
+ - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
167
+ - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
168
+ - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
211
169
 
212
- Parameters
213
- ----------
214
- cpu : int, default 1
215
- Number of CPUs required for this step.
216
- gpu : int, optional, default None
217
- Number of GPUs required for this step.
218
- disk : int, optional, default None
219
- Disk size (in MB) required for this step. Only applies on Kubernetes.
220
- memory : int, default 4096
221
- Memory size (in MB) required for this step.
222
- shared_memory : int, optional, default None
223
- The value for the size (in MiB) of the /dev/shm volume for this step.
224
- This parameter maps to the `--shm-size` option in Docker.
170
+ - If repo is found in the datastore:
171
+ - Loads it directly from datastore to local path (can be temporary directory or specified path)
225
172
  """
226
173
  ...
227
174
 
228
- 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]]]:
175
+ @typing.overload
176
+ 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]]]:
229
177
  """
230
- This decorator is used to run NIM containers in Metaflow tasks as sidecars.
178
+ Specifies a timeout for your step.
231
179
 
232
- User code call
233
- -----------
234
- @nim(
235
- models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
236
- backend='managed'
237
- )
180
+ This decorator is useful if this step may hang indefinitely.
238
181
 
239
- Valid backend options
240
- ---------------------
241
- - 'managed': Outerbounds selects a compute provider based on the model.
182
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
183
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
184
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
242
185
 
243
- Valid model options
244
- ----------------
245
- - 'meta/llama3-8b-instruct': 8B parameter model
246
- - 'meta/llama3-70b-instruct': 70B parameter model
247
- - any model here: https://nvcf.ngc.nvidia.com/functions?filter=nvidia-functions
186
+ Note that all the values specified in parameters are added together so if you specify
187
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
248
188
 
249
189
 
250
190
  Parameters
251
191
  ----------
252
- models: list[NIM]
253
- List of NIM containers running models in sidecars.
254
- backend: str
255
- Compute provider to run the NIM container.
256
- queue_timeout : int
257
- Time to keep the job in NVCF's queue.
192
+ seconds : int, default 0
193
+ Number of seconds to wait prior to timing out.
194
+ minutes : int, default 0
195
+ Number of minutes to wait prior to timing out.
196
+ hours : int, default 0
197
+ Number of hours to wait prior to timing out.
258
198
  """
259
199
  ...
260
200
 
261
201
  @typing.overload
262
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
263
- """
264
- Internal decorator to support Fast bakery
265
- Parameters
266
- ----------
267
- """
202
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
268
203
  ...
269
204
 
270
205
  @typing.overload
271
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
206
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
272
207
  ...
273
208
 
274
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
209
+ 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):
275
210
  """
276
- Internal decorator to support Fast bakery
211
+ Specifies a timeout for your step.
212
+
213
+ This decorator is useful if this step may hang indefinitely.
214
+
215
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
216
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
217
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
218
+
219
+ Note that all the values specified in parameters are added together so if you specify
220
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
221
+
222
+
277
223
  Parameters
278
224
  ----------
225
+ seconds : int, default 0
226
+ Number of seconds to wait prior to timing out.
227
+ minutes : int, default 0
228
+ Number of minutes to wait prior to timing out.
229
+ hours : int, default 0
230
+ Number of hours to wait prior to timing out.
279
231
  """
280
232
  ...
281
233
 
282
234
  @typing.overload
283
- 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]]]:
235
+ 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]]]:
284
236
  """
285
- Specifies the PyPI packages for the step.
237
+ Specifies that the step will success under all circumstances.
286
238
 
287
- Information in this decorator will augment any
288
- attributes set in the `@pyi_base` flow-level decorator. Hence,
289
- you can use `@pypi_base` to set packages required by all
290
- steps and use `@pypi` to specify step-specific overrides.
239
+ The decorator will create an optional artifact, specified by `var`, which
240
+ contains the exception raised. You can use it to detect the presence
241
+ of errors, indicating that all happy-path artifacts produced by the step
242
+ are missing.
291
243
 
292
244
 
293
245
  Parameters
294
246
  ----------
295
- packages : Dict[str, str], default: {}
296
- Packages to use for this step. The key is the name of the package
297
- and the value is the version to use.
298
- python : str, optional, default: None
299
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
300
- that the version used will correspond to the version of the Python interpreter used to start the run.
247
+ var : str, optional, default None
248
+ Name of the artifact in which to store the caught exception.
249
+ If not specified, the exception is not stored.
250
+ print_exception : bool, default True
251
+ Determines whether or not the exception is printed to
252
+ stdout when caught.
301
253
  """
302
254
  ...
303
255
 
304
256
  @typing.overload
305
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
257
+ def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
306
258
  ...
307
259
 
308
260
  @typing.overload
309
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
261
+ def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
310
262
  ...
311
263
 
312
- 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):
264
+ 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):
313
265
  """
314
- Specifies the PyPI packages for the step.
266
+ Specifies that the step will success under all circumstances.
315
267
 
316
- Information in this decorator will augment any
317
- attributes set in the `@pyi_base` flow-level decorator. Hence,
318
- you can use `@pypi_base` to set packages required by all
319
- steps and use `@pypi` to specify step-specific overrides.
268
+ The decorator will create an optional artifact, specified by `var`, which
269
+ contains the exception raised. You can use it to detect the presence
270
+ of errors, indicating that all happy-path artifacts produced by the step
271
+ are missing.
320
272
 
321
273
 
322
274
  Parameters
323
275
  ----------
324
- packages : Dict[str, str], default: {}
325
- Packages to use for this step. The key is the name of the package
326
- and the value is the version to use.
327
- python : str, optional, default: None
328
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
329
- that the version used will correspond to the version of the Python interpreter used to start the run.
276
+ var : str, optional, default None
277
+ Name of the artifact in which to store the caught exception.
278
+ If not specified, the exception is not stored.
279
+ print_exception : bool, default True
280
+ Determines whether or not the exception is printed to
281
+ stdout when caught.
330
282
  """
331
283
  ...
332
284
 
333
285
  @typing.overload
334
- 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 parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
335
287
  """
336
- Specifies secrets to be retrieved and injected as environment variables prior to
337
- the execution of a step.
338
-
339
-
340
288
  Parameters
341
289
  ----------
342
- sources : List[Union[str, Dict[str, Any]]], default: []
343
- List of secret specs, defining how the secrets are to be retrieved
344
290
  """
345
291
  ...
346
292
 
347
293
  @typing.overload
348
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
349
- ...
350
-
351
- @typing.overload
352
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
294
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
353
295
  ...
354
296
 
355
- 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]]] = []):
297
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
356
298
  """
357
- Specifies secrets to be retrieved and injected as environment variables prior to
358
- the execution of a step.
359
-
360
-
361
299
  Parameters
362
300
  ----------
363
- sources : List[Union[str, Dict[str, Any]]], default: []
364
- List of secret specs, defining how the secrets are to be retrieved
365
301
  """
366
302
  ...
367
303
 
368
304
  @typing.overload
369
- 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]]]:
305
+ 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]]]:
370
306
  """
371
- Enables checkpointing for a step.
307
+ Enables loading / saving of models within a step.
372
308
 
373
309
 
374
310
 
375
311
  Parameters
376
312
  ----------
377
- load_policy : str, default: "fresh"
378
- The policy for loading the checkpoint. The following policies are supported:
379
- - "eager": Loads the the latest available checkpoint within the namespace.
380
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
381
- will be loaded at the start of the task.
382
- - "none": Do not load any checkpoint
383
- - "fresh": Loads the lastest checkpoint created within the running Task.
384
- This mode helps loading checkpoints across various retry attempts of the same task.
385
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
386
- created within the task will be loaded when the task is retries execution on failure.
313
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
314
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
315
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
316
+ - `current.checkpoint`
317
+ - `current.model`
318
+ - `current.huggingface_hub`
319
+
320
+ 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
321
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
322
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
387
323
 
388
324
  temp_dir_root : str, default: None
389
- The root directory under which `current.checkpoint.directory` will be created.
325
+ The root directory under which `current.model.loaded` will store loaded models
390
326
  """
391
327
  ...
392
328
 
393
329
  @typing.overload
394
- def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
330
+ def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
395
331
  ...
396
332
 
397
333
  @typing.overload
398
- def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
334
+ def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
399
335
  ...
400
336
 
401
- 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):
337
+ 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):
402
338
  """
403
- Enables checkpointing for a step.
339
+ Enables loading / saving of models within a step.
404
340
 
405
341
 
406
342
 
407
343
  Parameters
408
344
  ----------
409
- load_policy : str, default: "fresh"
410
- The policy for loading the checkpoint. The following policies are supported:
411
- - "eager": Loads the the latest available checkpoint within the namespace.
412
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
413
- will be loaded at the start of the task.
414
- - "none": Do not load any checkpoint
415
- - "fresh": Loads the lastest checkpoint created within the running Task.
416
- This mode helps loading checkpoints across various retry attempts of the same task.
417
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
418
- created within the task will be loaded when the task is retries execution on failure.
345
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
346
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
347
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
348
+ - `current.checkpoint`
349
+ - `current.model`
350
+ - `current.huggingface_hub`
351
+
352
+ 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
353
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
354
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
419
355
 
420
356
  temp_dir_root : str, default: None
421
- The root directory under which `current.checkpoint.directory` will be created.
357
+ The root directory under which `current.model.loaded` will store loaded models
358
+ """
359
+ ...
360
+
361
+ 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]]]:
362
+ """
363
+ This decorator is used to run NIM containers in Metaflow tasks as sidecars.
364
+
365
+ User code call
366
+ -----------
367
+ @nim(
368
+ models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
369
+ backend='managed'
370
+ )
371
+
372
+ Valid backend options
373
+ ---------------------
374
+ - 'managed': Outerbounds selects a compute provider based on the model.
375
+
376
+ Valid model options
377
+ ----------------
378
+ - 'meta/llama3-8b-instruct': 8B parameter model
379
+ - 'meta/llama3-70b-instruct': 70B parameter model
380
+ - any model here: https://nvcf.ngc.nvidia.com/functions?filter=nvidia-functions
381
+
382
+
383
+ Parameters
384
+ ----------
385
+ models: list[NIM]
386
+ List of NIM containers running models in sidecars.
387
+ backend: str
388
+ Compute provider to run the NIM container.
389
+ queue_timeout : int
390
+ Time to keep the job in NVCF's queue.
422
391
  """
423
392
  ...
424
393
 
@@ -478,209 +447,145 @@ def retry(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
478
447
  ...
479
448
 
480
449
  @typing.overload
481
- 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]]]:
450
+ 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]]]:
482
451
  """
483
- Specifies environment variables to be set prior to the execution of a step.
452
+ Enables checkpointing for a step.
453
+
484
454
 
485
455
 
486
456
  Parameters
487
457
  ----------
488
- vars : Dict[str, str], default {}
489
- Dictionary of environment variables to set.
458
+ load_policy : str, default: "fresh"
459
+ The policy for loading the checkpoint. The following policies are supported:
460
+ - "eager": Loads the the latest available checkpoint within the namespace.
461
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
462
+ will be loaded at the start of the task.
463
+ - "none": Do not load any checkpoint
464
+ - "fresh": Loads the lastest checkpoint created within the running Task.
465
+ This mode helps loading checkpoints across various retry attempts of the same task.
466
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
467
+ created within the task will be loaded when the task is retries execution on failure.
468
+
469
+ temp_dir_root : str, default: None
470
+ The root directory under which `current.checkpoint.directory` will be created.
490
471
  """
491
472
  ...
492
473
 
493
474
  @typing.overload
494
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
475
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
495
476
  ...
496
477
 
497
478
  @typing.overload
498
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
479
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
499
480
  ...
500
481
 
501
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
482
+ 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):
502
483
  """
503
- Specifies environment variables to be set prior to the execution of a step.
504
-
484
+ Enables checkpointing for a step.
505
485
 
506
- Parameters
507
- ----------
508
- vars : Dict[str, str], default {}
509
- Dictionary of environment variables to set.
510
- """
511
- ...
512
-
513
- 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]]]:
514
- """
515
- Decorator that helps cache, version and store models/datasets from huggingface hub.
516
486
 
517
487
 
518
488
  Parameters
519
489
  ----------
520
- temp_dir_root : str, optional
521
- The root directory that will hold the temporary directory where objects will be downloaded.
522
-
523
- load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
524
- The list of repos (models/datasets) to load.
525
-
526
- Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
527
-
528
- - If repo (model/dataset) is not found in the datastore:
529
- - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
530
- - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
531
- - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
490
+ load_policy : str, default: "fresh"
491
+ The policy for loading the checkpoint. The following policies are supported:
492
+ - "eager": Loads the the latest available checkpoint within the namespace.
493
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
494
+ will be loaded at the start of the task.
495
+ - "none": Do not load any checkpoint
496
+ - "fresh": Loads the lastest checkpoint created within the running Task.
497
+ This mode helps loading checkpoints across various retry attempts of the same task.
498
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
499
+ created within the task will be loaded when the task is retries execution on failure.
532
500
 
533
- - If repo is found in the datastore:
534
- - Loads it directly from datastore to local path (can be temporary directory or specified path)
501
+ temp_dir_root : str, default: None
502
+ The root directory under which `current.checkpoint.directory` will be created.
535
503
  """
536
504
  ...
537
505
 
538
506
  @typing.overload
539
- 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]]]:
507
+ def card(*, type: str = 'default', id: typing.Optional[str] = None, options: typing.Dict[str, typing.Any] = {}, timeout: int = 45) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
540
508
  """
541
- Specifies the Conda environment for the step.
509
+ Creates a human-readable report, a Metaflow Card, after this step completes.
542
510
 
543
- Information in this decorator will augment any
544
- attributes set in the `@conda_base` flow-level decorator. Hence,
545
- you can use `@conda_base` to set packages required by all
546
- steps and use `@conda` to specify step-specific overrides.
511
+ Note that you may add multiple `@card` decorators in a step with different parameters.
547
512
 
548
513
 
549
514
  Parameters
550
515
  ----------
551
- packages : Dict[str, str], default {}
552
- Packages to use for this step. The key is the name of the package
553
- and the value is the version to use.
554
- libraries : Dict[str, str], default {}
555
- Supported for backward compatibility. When used with packages, packages will take precedence.
556
- python : str, optional, default None
557
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
558
- that the version used will correspond to the version of the Python interpreter used to start the run.
559
- disabled : bool, default False
560
- If set to True, disables @conda.
516
+ type : str, default 'default'
517
+ Card type.
518
+ id : str, optional, default None
519
+ If multiple cards are present, use this id to identify this card.
520
+ options : Dict[str, Any], default {}
521
+ Options passed to the card. The contents depend on the card type.
522
+ timeout : int, default 45
523
+ Interrupt reporting if it takes more than this many seconds.
561
524
  """
562
525
  ...
563
526
 
564
527
  @typing.overload
565
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
528
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
566
529
  ...
567
530
 
568
531
  @typing.overload
569
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
532
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
570
533
  ...
571
534
 
572
- 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):
535
+ def card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, type: str = 'default', id: typing.Optional[str] = None, options: typing.Dict[str, typing.Any] = {}, timeout: int = 45):
573
536
  """
574
- Specifies the Conda environment for the step.
537
+ Creates a human-readable report, a Metaflow Card, after this step completes.
575
538
 
576
- Information in this decorator will augment any
577
- attributes set in the `@conda_base` flow-level decorator. Hence,
578
- you can use `@conda_base` to set packages required by all
579
- steps and use `@conda` to specify step-specific overrides.
539
+ Note that you may add multiple `@card` decorators in a step with different parameters.
580
540
 
581
541
 
582
542
  Parameters
583
543
  ----------
584
- packages : Dict[str, str], default {}
585
- Packages to use for this step. The key is the name of the package
586
- and the value is the version to use.
587
- libraries : Dict[str, str], default {}
588
- Supported for backward compatibility. When used with packages, packages will take precedence.
589
- python : str, optional, default None
590
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
591
- that the version used will correspond to the version of the Python interpreter used to start the run.
592
- disabled : bool, default False
593
- If set to True, disables @conda.
544
+ type : str, default 'default'
545
+ Card type.
546
+ id : str, optional, default None
547
+ If multiple cards are present, use this id to identify this card.
548
+ options : Dict[str, Any], default {}
549
+ Options passed to the card. The contents depend on the card type.
550
+ timeout : int, default 45
551
+ Interrupt reporting if it takes more than this many seconds.
594
552
  """
595
553
  ...
596
554
 
597
555
  @typing.overload
598
- 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]]]:
556
+ 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]]]:
599
557
  """
600
- Specifies that the step will success under all circumstances.
601
-
602
- The decorator will create an optional artifact, specified by `var`, which
603
- contains the exception raised. You can use it to detect the presence
604
- of errors, indicating that all happy-path artifacts produced by the step
605
- are missing.
558
+ Specifies environment variables to be set prior to the execution of a step.
606
559
 
607
560
 
608
561
  Parameters
609
562
  ----------
610
- var : str, optional, default None
611
- Name of the artifact in which to store the caught exception.
612
- If not specified, the exception is not stored.
613
- print_exception : bool, default True
614
- Determines whether or not the exception is printed to
615
- stdout when caught.
563
+ vars : Dict[str, str], default {}
564
+ Dictionary of environment variables to set.
616
565
  """
617
566
  ...
618
567
 
619
568
  @typing.overload
620
- def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
569
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
621
570
  ...
622
571
 
623
572
  @typing.overload
624
- def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
625
- ...
626
-
627
- 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):
628
- """
629
- Specifies that the step will success under all circumstances.
630
-
631
- The decorator will create an optional artifact, specified by `var`, which
632
- contains the exception raised. You can use it to detect the presence
633
- of errors, indicating that all happy-path artifacts produced by the step
634
- are missing.
635
-
636
-
637
- Parameters
638
- ----------
639
- var : str, optional, default None
640
- Name of the artifact in which to store the caught exception.
641
- If not specified, the exception is not stored.
642
- print_exception : bool, default True
643
- Determines whether or not the exception is printed to
644
- stdout when caught.
645
- """
573
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
646
574
  ...
647
575
 
648
- 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]]]:
576
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
649
577
  """
650
- Specifies that this step should execute on DGX cloud.
578
+ Specifies environment variables to be set prior to the execution of a step.
651
579
 
652
580
 
653
581
  Parameters
654
582
  ----------
655
- gpu : int
656
- Number of GPUs to use.
657
- gpu_type : str
658
- Type of Nvidia GPU to use.
659
- queue_timeout : int
660
- Time to keep the job in NVCF's queue.
661
- """
662
- ...
663
-
664
- @typing.overload
665
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
666
- """
667
- Parameters
668
- ----------
669
- """
670
- ...
671
-
672
- @typing.overload
673
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
674
- ...
675
-
676
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
677
- """
678
- Parameters
679
- ----------
583
+ vars : Dict[str, str], default {}
584
+ Dictionary of environment variables to set.
680
585
  """
681
586
  ...
682
587
 
683
- def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: typing.Optional[str] = None, image_pull_policy: str = 'KUBERNETES_IMAGE_PULL_POLICY', service_account: str = 'METAFLOW_KUBERNETES_SERVICE_ACCOUNT', secrets: typing.Optional[typing.List[str]] = None, node_selector: typing.Union[typing.Dict[str, str], str, None] = None, namespace: str = 'METAFLOW_KUBERNETES_NAMESPACE', gpu: typing.Optional[int] = None, gpu_vendor: str = 'KUBERNETES_GPU_VENDOR', tolerations: typing.List[str] = [], labels: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_LABELS', annotations: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_ANNOTATIONS', use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = '/metaflow_temp', persistent_volume_claims: typing.Optional[typing.Dict[str, str]] = None, shared_memory: typing.Optional[int] = None, port: typing.Optional[int] = None, compute_pool: typing.Optional[str] = None, hostname_resolution_timeout: int = 600, qos: str = 'Burstable') -> 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]]]:
588
+ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: typing.Optional[str] = None, image_pull_policy: str = 'KUBERNETES_IMAGE_PULL_POLICY', service_account: str = 'METAFLOW_KUBERNETES_SERVICE_ACCOUNT', secrets: typing.Optional[typing.List[str]] = None, node_selector: typing.Union[typing.Dict[str, str], str, None] = None, namespace: str = 'METAFLOW_KUBERNETES_NAMESPACE', gpu: typing.Optional[int] = None, gpu_vendor: str = 'KUBERNETES_GPU_VENDOR', tolerations: typing.List[str] = [], labels: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_LABELS', annotations: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_ANNOTATIONS', use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = '/metaflow_temp', persistent_volume_claims: typing.Optional[typing.Dict[str, str]] = None, shared_memory: typing.Optional[int] = None, port: typing.Optional[int] = None, compute_pool: typing.Optional[str] = None, hostname_resolution_timeout: int = 600, qos: str = 'Burstable', extended_resources: typing.Optional[typing.Dict[str, 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]]]:
684
589
  """
685
590
  Specifies that this step should execute on Kubernetes.
686
591
 
@@ -754,171 +659,321 @@ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: ty
754
659
  Only applicable when @parallel is used.
755
660
  qos: str, default: Burstable
756
661
  Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
662
+ extended_resources: Dict[str, str], optional, default None
663
+ Extended resources to be requested for the pod.
664
+ https://kubernetes.io/docs/tasks/administer-cluster/extended-resource-node/
757
665
  """
758
666
  ...
759
667
 
760
668
  @typing.overload
761
- 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]]]:
669
+ 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]]]:
762
670
  """
763
- Specifies a timeout for your step.
764
-
765
- This decorator is useful if this step may hang indefinitely.
766
-
767
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
768
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
769
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
671
+ Specifies the Conda environment for the step.
770
672
 
771
- Note that all the values specified in parameters are added together so if you specify
772
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
673
+ Information in this decorator will augment any
674
+ attributes set in the `@conda_base` flow-level decorator. Hence,
675
+ you can use `@conda_base` to set packages required by all
676
+ steps and use `@conda` to specify step-specific overrides.
773
677
 
774
678
 
775
679
  Parameters
776
680
  ----------
777
- seconds : int, default 0
778
- Number of seconds to wait prior to timing out.
779
- minutes : int, default 0
780
- Number of minutes to wait prior to timing out.
781
- hours : int, default 0
782
- Number of hours to wait prior to timing out.
681
+ packages : Dict[str, str], default {}
682
+ Packages to use for this step. The key is the name of the package
683
+ and the value is the version to use.
684
+ libraries : Dict[str, str], default {}
685
+ Supported for backward compatibility. When used with packages, packages will take precedence.
686
+ python : str, optional, default None
687
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
688
+ that the version used will correspond to the version of the Python interpreter used to start the run.
689
+ disabled : bool, default False
690
+ If set to True, disables @conda.
783
691
  """
784
692
  ...
785
693
 
786
694
  @typing.overload
787
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
695
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
788
696
  ...
789
697
 
790
698
  @typing.overload
791
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
699
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
792
700
  ...
793
701
 
794
- 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):
702
+ 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):
795
703
  """
796
- Specifies a timeout for your step.
704
+ Specifies the Conda environment for the step.
797
705
 
798
- This decorator is useful if this step may hang indefinitely.
706
+ Information in this decorator will augment any
707
+ attributes set in the `@conda_base` flow-level decorator. Hence,
708
+ you can use `@conda_base` to set packages required by all
709
+ steps and use `@conda` to specify step-specific overrides.
799
710
 
800
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
801
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
802
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
803
711
 
804
- Note that all the values specified in parameters are added together so if you specify
805
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
712
+ Parameters
713
+ ----------
714
+ packages : Dict[str, str], default {}
715
+ Packages to use for this step. The key is the name of the package
716
+ and the value is the version to use.
717
+ libraries : Dict[str, str], default {}
718
+ Supported for backward compatibility. When used with packages, packages will take precedence.
719
+ python : str, optional, default None
720
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
721
+ that the version used will correspond to the version of the Python interpreter used to start the run.
722
+ disabled : bool, default False
723
+ If set to True, disables @conda.
724
+ """
725
+ ...
726
+
727
+ @typing.overload
728
+ 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]]]:
729
+ """
730
+ Specifies the resources needed when executing this step.
731
+
732
+ Use `@resources` to specify the resource requirements
733
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
734
+
735
+ You can choose the compute layer on the command line by executing e.g.
736
+ ```
737
+ python myflow.py run --with batch
738
+ ```
739
+ or
740
+ ```
741
+ python myflow.py run --with kubernetes
742
+ ```
743
+ which executes the flow on the desired system using the
744
+ requirements specified in `@resources`.
806
745
 
807
746
 
808
747
  Parameters
809
748
  ----------
810
- seconds : int, default 0
811
- Number of seconds to wait prior to timing out.
812
- minutes : int, default 0
813
- Number of minutes to wait prior to timing out.
814
- hours : int, default 0
815
- Number of hours to wait prior to timing out.
749
+ cpu : int, default 1
750
+ Number of CPUs required for this step.
751
+ gpu : int, optional, default None
752
+ Number of GPUs required for this step.
753
+ disk : int, optional, default None
754
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
755
+ memory : int, default 4096
756
+ Memory size (in MB) required for this step.
757
+ shared_memory : int, optional, default None
758
+ The value for the size (in MiB) of the /dev/shm volume for this step.
759
+ This parameter maps to the `--shm-size` option in Docker.
816
760
  """
817
761
  ...
818
762
 
819
763
  @typing.overload
820
- 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]]]:
764
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
765
+ ...
766
+
767
+ @typing.overload
768
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
769
+ ...
770
+
771
+ 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):
821
772
  """
822
- Enables loading / saving of models within a step.
773
+ Specifies the resources needed when executing this step.
823
774
 
775
+ Use `@resources` to specify the resource requirements
776
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
777
+
778
+ You can choose the compute layer on the command line by executing e.g.
779
+ ```
780
+ python myflow.py run --with batch
781
+ ```
782
+ or
783
+ ```
784
+ python myflow.py run --with kubernetes
785
+ ```
786
+ which executes the flow on the desired system using the
787
+ requirements specified in `@resources`.
824
788
 
825
789
 
826
790
  Parameters
827
791
  ----------
828
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
829
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
830
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
831
- - `current.checkpoint`
832
- - `current.model`
833
- - `current.huggingface_hub`
792
+ cpu : int, default 1
793
+ Number of CPUs required for this step.
794
+ gpu : int, optional, default None
795
+ Number of GPUs required for this step.
796
+ disk : int, optional, default None
797
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
798
+ memory : int, default 4096
799
+ Memory size (in MB) required for this step.
800
+ shared_memory : int, optional, default None
801
+ The value for the size (in MiB) of the /dev/shm volume for this step.
802
+ This parameter maps to the `--shm-size` option in Docker.
803
+ """
804
+ ...
805
+
806
+ @typing.overload
807
+ 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]]]:
808
+ """
809
+ Specifies secrets to be retrieved and injected as environment variables prior to
810
+ the execution of a step.
834
811
 
835
- 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
836
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
837
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
838
812
 
839
- temp_dir_root : str, default: None
840
- The root directory under which `current.model.loaded` will store loaded models
813
+ Parameters
814
+ ----------
815
+ sources : List[Union[str, Dict[str, Any]]], default: []
816
+ List of secret specs, defining how the secrets are to be retrieved
841
817
  """
842
818
  ...
843
819
 
844
820
  @typing.overload
845
- def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
821
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
846
822
  ...
847
823
 
848
824
  @typing.overload
849
- def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
825
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
850
826
  ...
851
827
 
852
- 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):
828
+ 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]]] = []):
853
829
  """
854
- Enables loading / saving of models within a step.
830
+ Specifies secrets to be retrieved and injected as environment variables prior to
831
+ the execution of a step.
855
832
 
856
833
 
834
+ Parameters
835
+ ----------
836
+ sources : List[Union[str, Dict[str, Any]]], default: []
837
+ List of secret specs, defining how the secrets are to be retrieved
838
+ """
839
+ ...
840
+
841
+ 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]]]:
842
+ """
843
+ Specifies that this step should execute on DGX cloud.
844
+
857
845
 
858
846
  Parameters
859
847
  ----------
860
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
861
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
862
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
863
- - `current.checkpoint`
864
- - `current.model`
865
- - `current.huggingface_hub`
848
+ gpu : int
849
+ Number of GPUs to use.
850
+ gpu_type : str
851
+ Type of Nvidia GPU to use.
852
+ queue_timeout : int
853
+ Time to keep the job in NVCF's queue.
854
+ """
855
+ ...
856
+
857
+ @typing.overload
858
+ 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]]]:
859
+ """
860
+ Specifies the PyPI packages for the step.
866
861
 
867
- 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
868
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
869
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
862
+ Information in this decorator will augment any
863
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
864
+ you can use `@pypi_base` to set packages required by all
865
+ steps and use `@pypi` to specify step-specific overrides.
870
866
 
871
- temp_dir_root : str, default: None
872
- The root directory under which `current.model.loaded` will store loaded models
867
+
868
+ Parameters
869
+ ----------
870
+ packages : Dict[str, str], default: {}
871
+ Packages to use for this step. The key is the name of the package
872
+ and the value is the version to use.
873
+ python : str, optional, default: None
874
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
875
+ that the version used will correspond to the version of the Python interpreter used to start the run.
873
876
  """
874
877
  ...
875
878
 
876
879
  @typing.overload
877
- def card(*, type: str = 'default', id: typing.Optional[str] = None, options: typing.Dict[str, typing.Any] = {}, timeout: int = 45) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
880
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
881
+ ...
882
+
883
+ @typing.overload
884
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
885
+ ...
886
+
887
+ 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):
878
888
  """
879
- Creates a human-readable report, a Metaflow Card, after this step completes.
889
+ Specifies the PyPI packages for the step.
890
+
891
+ Information in this decorator will augment any
892
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
893
+ you can use `@pypi_base` to set packages required by all
894
+ steps and use `@pypi` to specify step-specific overrides.
895
+
896
+
897
+ Parameters
898
+ ----------
899
+ packages : Dict[str, str], default: {}
900
+ Packages to use for this step. The key is the name of the package
901
+ and the value is the version to use.
902
+ python : str, optional, default: None
903
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
904
+ that the version used will correspond to the version of the Python interpreter used to start the run.
905
+ """
906
+ ...
907
+
908
+ @typing.overload
909
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
910
+ """
911
+ Internal decorator to support Fast bakery
912
+ Parameters
913
+ ----------
914
+ """
915
+ ...
916
+
917
+ @typing.overload
918
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
919
+ ...
920
+
921
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
922
+ """
923
+ Internal decorator to support Fast bakery
924
+ Parameters
925
+ ----------
926
+ """
927
+ ...
928
+
929
+ @typing.overload
930
+ 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]]:
931
+ """
932
+ Specifies the Conda environment for all steps of the flow.
880
933
 
881
- Note that you may add multiple `@card` decorators in a step with different parameters.
934
+ Use `@conda_base` to set common libraries required by all
935
+ steps and use `@conda` to specify step-specific additions.
882
936
 
883
937
 
884
938
  Parameters
885
939
  ----------
886
- type : str, default 'default'
887
- Card type.
888
- id : str, optional, default None
889
- If multiple cards are present, use this id to identify this card.
890
- options : Dict[str, Any], default {}
891
- Options passed to the card. The contents depend on the card type.
892
- timeout : int, default 45
893
- Interrupt reporting if it takes more than this many seconds.
940
+ packages : Dict[str, str], default {}
941
+ Packages to use for this flow. The key is the name of the package
942
+ and the value is the version to use.
943
+ libraries : Dict[str, str], default {}
944
+ Supported for backward compatibility. When used with packages, packages will take precedence.
945
+ python : str, optional, default None
946
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
947
+ that the version used will correspond to the version of the Python interpreter used to start the run.
948
+ disabled : bool, default False
949
+ If set to True, disables Conda.
894
950
  """
895
951
  ...
896
952
 
897
953
  @typing.overload
898
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
899
- ...
900
-
901
- @typing.overload
902
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
954
+ def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
903
955
  ...
904
956
 
905
- def card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, type: str = 'default', id: typing.Optional[str] = None, options: typing.Dict[str, typing.Any] = {}, timeout: int = 45):
957
+ 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):
906
958
  """
907
- Creates a human-readable report, a Metaflow Card, after this step completes.
959
+ Specifies the Conda environment for all steps of the flow.
908
960
 
909
- Note that you may add multiple `@card` decorators in a step with different parameters.
961
+ Use `@conda_base` to set common libraries required by all
962
+ steps and use `@conda` to specify step-specific additions.
910
963
 
911
964
 
912
965
  Parameters
913
966
  ----------
914
- type : str, default 'default'
915
- Card type.
916
- id : str, optional, default None
917
- If multiple cards are present, use this id to identify this card.
918
- options : Dict[str, Any], default {}
919
- Options passed to the card. The contents depend on the card type.
920
- timeout : int, default 45
921
- Interrupt reporting if it takes more than this many seconds.
967
+ packages : Dict[str, str], default {}
968
+ Packages to use for this flow. The key is the name of the package
969
+ and the value is the version to use.
970
+ libraries : Dict[str, str], default {}
971
+ Supported for backward compatibility. When used with packages, packages will take precedence.
972
+ python : str, optional, default None
973
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
974
+ that the version used will correspond to the version of the Python interpreter used to start the run.
975
+ disabled : bool, default False
976
+ If set to True, disables Conda.
922
977
  """
923
978
  ...
924
979
 
@@ -965,89 +1020,114 @@ def airflow_external_task_sensor(*, timeout: int, poke_interval: int, mode: str,
965
1020
  """
966
1021
  ...
967
1022
 
968
- @typing.overload
969
- 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]]:
1023
+ def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None):
970
1024
  """
971
- Specifies the Conda environment for all steps of the flow.
1025
+ Allows setting external datastores to save data for the
1026
+ `@checkpoint`/`@model`/`@huggingface_hub` decorators. This is useful
1027
+ when the compute medium lives in a different geographical location
1028
+ than metaflow's datastore (e.g. S3, GCS, Azure Blob Storage).
972
1029
 
973
- Use `@conda_base` to set common libraries required by all
974
- steps and use `@conda` to specify step-specific additions.
1030
+ Parameters:
1031
+ ----------
1032
+
1033
+ type: str
1034
+ The type of the datastore. Can be one of 's3', 'gcs', 'azure' or any other supported metaflow Datastore.
1035
+
1036
+ config: dict or Callable
1037
+ Dictionary of configuration options for the datastore. The following keys are required:
1038
+ - root: The root path in the datastore where the data will be saved. (needs to be in the format expected by the datastore)
1039
+ - example: 's3://bucket-name/path/to/root'
1040
+ - example: 'gs://bucket-name/path/to/root'
1041
+ - example: 'https://myblockacc.blob.core.windows.net/metaflow/'
1042
+ - role_arn (optional): AWS IAM role to access s3 bucket (only when `type` is 's3')
1043
+ - session_vars (optional): AWS session variables to access s3 bucket (only when `type` is 's3')
1044
+ - client_params (optional): AWS client parameters to access s3 bucket (only when `type` is 's3')
1045
+ Parameters
1046
+ ----------
1047
+ """
1048
+ ...
1049
+
1050
+ @typing.overload
1051
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1052
+ """
1053
+ Specifies the PyPI packages for all steps of the flow.
975
1054
 
1055
+ Use `@pypi_base` to set common packages required by all
1056
+ steps and use `@pypi` to specify step-specific overrides.
976
1057
 
977
1058
  Parameters
978
1059
  ----------
979
- packages : Dict[str, str], default {}
1060
+ packages : Dict[str, str], default: {}
980
1061
  Packages to use for this flow. The key is the name of the package
981
1062
  and the value is the version to use.
982
- libraries : Dict[str, str], default {}
983
- Supported for backward compatibility. When used with packages, packages will take precedence.
984
- python : str, optional, default None
1063
+ python : str, optional, default: None
985
1064
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
986
1065
  that the version used will correspond to the version of the Python interpreter used to start the run.
987
- disabled : bool, default False
988
- If set to True, disables Conda.
989
1066
  """
990
1067
  ...
991
1068
 
992
1069
  @typing.overload
993
- def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1070
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
994
1071
  ...
995
1072
 
996
- 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):
1073
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
997
1074
  """
998
- Specifies the Conda environment for all steps of the flow.
999
-
1000
- Use `@conda_base` to set common libraries required by all
1001
- steps and use `@conda` to specify step-specific additions.
1075
+ Specifies the PyPI packages for all steps of the flow.
1002
1076
 
1077
+ Use `@pypi_base` to set common packages required by all
1078
+ steps and use `@pypi` to specify step-specific overrides.
1003
1079
 
1004
1080
  Parameters
1005
1081
  ----------
1006
- packages : Dict[str, str], default {}
1082
+ packages : Dict[str, str], default: {}
1007
1083
  Packages to use for this flow. The key is the name of the package
1008
1084
  and the value is the version to use.
1009
- libraries : Dict[str, str], default {}
1010
- Supported for backward compatibility. When used with packages, packages will take precedence.
1011
- python : str, optional, default None
1085
+ python : str, optional, default: None
1012
1086
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
1013
1087
  that the version used will correspond to the version of the Python interpreter used to start the run.
1014
- disabled : bool, default False
1015
- If set to True, disables Conda.
1016
1088
  """
1017
1089
  ...
1018
1090
 
1019
- def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1091
+ def airflow_s3_key_sensor(*, timeout: int, poke_interval: int, mode: str, exponential_backoff: bool, pool: str, soft_fail: bool, name: str, description: str, bucket_key: typing.Union[str, typing.List[str]], bucket_name: str, wildcard_match: bool, aws_conn_id: str, verify: bool) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1020
1092
  """
1021
- Specifies what flows belong to the same project.
1022
-
1023
- A project-specific namespace is created for all flows that
1024
- use the same `@project(name)`.
1093
+ The `@airflow_s3_key_sensor` decorator attaches a Airflow [S3KeySensor](https://airflow.apache.org/docs/apache-airflow-providers-amazon/stable/_api/airflow/providers/amazon/aws/sensors/s3/index.html#airflow.providers.amazon.aws.sensors.s3.S3KeySensor)
1094
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1095
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1096
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
1097
+ starts only after all sensors finish.
1025
1098
 
1026
1099
 
1027
1100
  Parameters
1028
1101
  ----------
1102
+ timeout : int
1103
+ Time, in seconds before the task times out and fails. (Default: 3600)
1104
+ poke_interval : int
1105
+ Time in seconds that the job should wait in between each try. (Default: 60)
1106
+ mode : str
1107
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1108
+ exponential_backoff : bool
1109
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1110
+ pool : str
1111
+ the slot pool this task should run in,
1112
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1113
+ soft_fail : bool
1114
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1029
1115
  name : str
1030
- Project name. Make sure that the name is unique amongst all
1031
- projects that use the same production scheduler. The name may
1032
- contain only lowercase alphanumeric characters and underscores.
1033
-
1034
- branch : Optional[str], default None
1035
- The branch to use. If not specified, the branch is set to
1036
- `user.<username>` unless `production` is set to `True`. This can
1037
- also be set on the command line using `--branch` as a top-level option.
1038
- It is an error to specify `branch` in the decorator and on the command line.
1039
-
1040
- production : bool, default False
1041
- Whether or not the branch is the production branch. This can also be set on the
1042
- command line using `--production` as a top-level option. It is an error to specify
1043
- `production` in the decorator and on the command line.
1044
- The project branch name will be:
1045
- - if `branch` is specified:
1046
- - if `production` is True: `prod.<branch>`
1047
- - if `production` is False: `test.<branch>`
1048
- - if `branch` is not specified:
1049
- - if `production` is True: `prod`
1050
- - if `production` is False: `user.<username>`
1116
+ Name of the sensor on Airflow
1117
+ description : str
1118
+ Description of sensor in the Airflow UI
1119
+ bucket_key : Union[str, List[str]]
1120
+ The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1121
+ When it's specified as a full s3:// url, please leave `bucket_name` as None
1122
+ bucket_name : str
1123
+ Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1124
+ When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1125
+ wildcard_match : bool
1126
+ whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1127
+ aws_conn_id : str
1128
+ a reference to the s3 connection on Airflow. (Default: None)
1129
+ verify : bool
1130
+ Whether or not to verify SSL certificates for S3 connection. (Default: None)
1051
1131
  """
1052
1132
  ...
1053
1133
 
@@ -1144,44 +1224,38 @@ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: t
1144
1224
  """
1145
1225
  ...
1146
1226
 
1147
- @typing.overload
1148
- def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1227
+ def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1149
1228
  """
1150
- Specifies the PyPI packages for all steps of the flow.
1229
+ Specifies what flows belong to the same project.
1230
+
1231
+ A project-specific namespace is created for all flows that
1232
+ use the same `@project(name)`.
1151
1233
 
1152
- Use `@pypi_base` to set common packages required by all
1153
- steps and use `@pypi` to specify step-specific overrides.
1154
1234
 
1155
1235
  Parameters
1156
1236
  ----------
1157
- packages : Dict[str, str], default: {}
1158
- Packages to use for this flow. The key is the name of the package
1159
- and the value is the version to use.
1160
- python : str, optional, default: None
1161
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1162
- that the version used will correspond to the version of the Python interpreter used to start the run.
1163
- """
1164
- ...
1165
-
1166
- @typing.overload
1167
- def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1168
- ...
1169
-
1170
- def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1171
- """
1172
- Specifies the PyPI packages for all steps of the flow.
1237
+ name : str
1238
+ Project name. Make sure that the name is unique amongst all
1239
+ projects that use the same production scheduler. The name may
1240
+ contain only lowercase alphanumeric characters and underscores.
1173
1241
 
1174
- Use `@pypi_base` to set common packages required by all
1175
- steps and use `@pypi` to specify step-specific overrides.
1242
+ branch : Optional[str], default None
1243
+ The branch to use. If not specified, the branch is set to
1244
+ `user.<username>` unless `production` is set to `True`. This can
1245
+ also be set on the command line using `--branch` as a top-level option.
1246
+ It is an error to specify `branch` in the decorator and on the command line.
1176
1247
 
1177
- Parameters
1178
- ----------
1179
- packages : Dict[str, str], default: {}
1180
- Packages to use for this flow. The key is the name of the package
1181
- and the value is the version to use.
1182
- python : str, optional, default: None
1183
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1184
- that the version used will correspond to the version of the Python interpreter used to start the run.
1248
+ production : bool, default False
1249
+ Whether or not the branch is the production branch. This can also be set on the
1250
+ command line using `--production` as a top-level option. It is an error to specify
1251
+ `production` in the decorator and on the command line.
1252
+ The project branch name will be:
1253
+ - if `branch` is specified:
1254
+ - if `production` is True: `prod.<branch>`
1255
+ - if `production` is False: `test.<branch>`
1256
+ - if `branch` is not specified:
1257
+ - if `production` is True: `prod`
1258
+ - if `production` is False: `user.<username>`
1185
1259
  """
1186
1260
  ...
1187
1261
 
@@ -1337,48 +1411,5 @@ def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *
1337
1411
  """
1338
1412
  ...
1339
1413
 
1340
- def airflow_s3_key_sensor(*, timeout: int, poke_interval: int, mode: str, exponential_backoff: bool, pool: str, soft_fail: bool, name: str, description: str, bucket_key: typing.Union[str, typing.List[str]], bucket_name: str, wildcard_match: bool, aws_conn_id: str, verify: bool) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1341
- """
1342
- The `@airflow_s3_key_sensor` decorator attaches a Airflow [S3KeySensor](https://airflow.apache.org/docs/apache-airflow-providers-amazon/stable/_api/airflow/providers/amazon/aws/sensors/s3/index.html#airflow.providers.amazon.aws.sensors.s3.S3KeySensor)
1343
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1344
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1345
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
1346
- starts only after all sensors finish.
1347
-
1348
-
1349
- Parameters
1350
- ----------
1351
- timeout : int
1352
- Time, in seconds before the task times out and fails. (Default: 3600)
1353
- poke_interval : int
1354
- Time in seconds that the job should wait in between each try. (Default: 60)
1355
- mode : str
1356
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1357
- exponential_backoff : bool
1358
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1359
- pool : str
1360
- the slot pool this task should run in,
1361
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1362
- soft_fail : bool
1363
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1364
- name : str
1365
- Name of the sensor on Airflow
1366
- description : str
1367
- Description of sensor in the Airflow UI
1368
- bucket_key : Union[str, List[str]]
1369
- The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1370
- When it's specified as a full s3:// url, please leave `bucket_name` as None
1371
- bucket_name : str
1372
- Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1373
- When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1374
- wildcard_match : bool
1375
- whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1376
- aws_conn_id : str
1377
- a reference to the s3 connection on Airflow. (Default: None)
1378
- verify : bool
1379
- Whether or not to verify SSL certificates for S3 connection. (Default: None)
1380
- """
1381
- ...
1382
-
1383
1414
  pkg_name: str
1384
1415