ob-metaflow-stubs 6.0.3.175rc1__py2.py3-none-any.whl → 6.0.3.176rc1__py2.py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (218) hide show
  1. metaflow-stubs/__init__.pyi +684 -678
  2. metaflow-stubs/cards.pyi +1 -1
  3. metaflow-stubs/cli.pyi +1 -1
  4. metaflow-stubs/cli_components/__init__.pyi +1 -1
  5. metaflow-stubs/cli_components/utils.pyi +1 -1
  6. metaflow-stubs/client/__init__.pyi +1 -1
  7. metaflow-stubs/client/core.pyi +5 -5
  8. metaflow-stubs/client/filecache.pyi +2 -2
  9. metaflow-stubs/events.pyi +2 -2
  10. metaflow-stubs/exception.pyi +1 -1
  11. metaflow-stubs/flowspec.pyi +4 -4
  12. metaflow-stubs/generated_for.txt +1 -1
  13. metaflow-stubs/includefile.pyi +3 -3
  14. metaflow-stubs/info_file.pyi +1 -1
  15. metaflow-stubs/metadata_provider/__init__.pyi +1 -1
  16. metaflow-stubs/metadata_provider/heartbeat.pyi +1 -1
  17. metaflow-stubs/metadata_provider/metadata.pyi +2 -2
  18. metaflow-stubs/metadata_provider/util.pyi +1 -1
  19. metaflow-stubs/metaflow_config.pyi +1 -1
  20. metaflow-stubs/metaflow_current.pyi +128 -128
  21. metaflow-stubs/metaflow_git.pyi +1 -1
  22. metaflow-stubs/mf_extensions/__init__.pyi +1 -1
  23. metaflow-stubs/mf_extensions/obcheckpoint/__init__.pyi +1 -1
  24. metaflow-stubs/mf_extensions/obcheckpoint/plugins/__init__.pyi +1 -1
  25. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/__init__.pyi +1 -1
  26. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/__init__.pyi +1 -1
  27. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/async_cards.pyi +1 -1
  28. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/deco_injection_mixin.pyi +1 -1
  29. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/extra_components.pyi +3 -3
  30. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/__init__.pyi +1 -1
  31. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/__init__.pyi +1 -1
  32. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/checkpoint_lister.pyi +4 -4
  33. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/lineage_card.pyi +1 -1
  34. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/checkpoint_storage.pyi +4 -4
  35. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/constructors.pyi +1 -1
  36. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/core.pyi +3 -3
  37. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/decorator.pyi +4 -4
  38. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/exceptions.pyi +1 -1
  39. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/final_api.pyi +2 -2
  40. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/lineage.pyi +1 -1
  41. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/__init__.pyi +1 -1
  42. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/context.pyi +2 -2
  43. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/core.pyi +1 -1
  44. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/decorator.pyi +1 -1
  45. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/exceptions.pyi +1 -1
  46. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/task_utils.pyi +2 -2
  47. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/utils.pyi +1 -1
  48. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastructures.pyi +2 -2
  49. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/exceptions.pyi +1 -1
  50. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/__init__.pyi +1 -1
  51. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/decorator.pyi +1 -1
  52. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/__init__.pyi +1 -1
  53. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/core.pyi +3 -3
  54. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/exceptions.pyi +1 -1
  55. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/model_storage.pyi +3 -3
  56. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/__init__.pyi +1 -1
  57. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/flowspec_utils.pyi +1 -1
  58. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/general.pyi +1 -1
  59. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/identity_utils.pyi +1 -1
  60. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/__init__.pyi +1 -1
  61. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/base.pyi +1 -1
  62. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/tar.pyi +2 -2
  63. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/tar_utils.pyi +2 -2
  64. metaflow-stubs/mf_extensions/outerbounds/__init__.pyi +1 -1
  65. metaflow-stubs/mf_extensions/outerbounds/plugins/__init__.pyi +1 -1
  66. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/__init__.pyi +1 -1
  67. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/async_cards.pyi +60 -0
  68. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/injector.pyi +1 -1
  69. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/__init__.pyi +1 -1
  70. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/coreweave.pyi +1 -1
  71. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/nebius.pyi +1 -1
  72. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/__init__.pyi +1 -1
  73. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/baker.pyi +4 -4
  74. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/docker_environment.pyi +2 -2
  75. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/fast_bakery.pyi +1 -1
  76. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/__init__.pyi +1 -1
  77. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/pod_killer.pyi +1 -1
  78. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/__init__.pyi +1 -1
  79. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/constants.pyi +1 -1
  80. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/exceptions.pyi +1 -1
  81. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/ollama.pyi +60 -2
  82. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/status_card.pyi +73 -0
  83. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/__init__.pyi +1 -1
  84. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/snowflake.pyi +1 -1
  85. metaflow-stubs/mf_extensions/outerbounds/profilers/__init__.pyi +1 -1
  86. metaflow-stubs/mf_extensions/outerbounds/profilers/gpu.pyi +1 -1
  87. metaflow-stubs/mf_extensions/outerbounds/remote_config.pyi +1 -1
  88. metaflow-stubs/mf_extensions/outerbounds/toplevel/__init__.pyi +1 -1
  89. metaflow-stubs/mf_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.pyi +1 -1
  90. metaflow-stubs/multicore_utils.pyi +1 -1
  91. metaflow-stubs/ob_internal.pyi +1 -1
  92. metaflow-stubs/parameters.pyi +3 -3
  93. metaflow-stubs/plugins/__init__.pyi +10 -10
  94. metaflow-stubs/plugins/airflow/__init__.pyi +1 -1
  95. metaflow-stubs/plugins/airflow/airflow_utils.pyi +1 -1
  96. metaflow-stubs/plugins/airflow/exception.pyi +1 -1
  97. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +1 -1
  98. metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi +1 -1
  99. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +1 -1
  100. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +1 -1
  101. metaflow-stubs/plugins/argo/__init__.pyi +1 -1
  102. metaflow-stubs/plugins/argo/argo_client.pyi +1 -1
  103. metaflow-stubs/plugins/argo/argo_events.pyi +1 -1
  104. metaflow-stubs/plugins/argo/argo_workflows.pyi +3 -3
  105. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +3 -3
  106. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +3 -3
  107. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +3 -3
  108. metaflow-stubs/plugins/aws/__init__.pyi +1 -1
  109. metaflow-stubs/plugins/aws/aws_client.pyi +1 -1
  110. metaflow-stubs/plugins/aws/aws_utils.pyi +1 -1
  111. metaflow-stubs/plugins/aws/batch/__init__.pyi +1 -1
  112. metaflow-stubs/plugins/aws/batch/batch.pyi +1 -1
  113. metaflow-stubs/plugins/aws/batch/batch_client.pyi +1 -1
  114. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +1 -1
  115. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +1 -1
  116. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +2 -2
  117. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +1 -1
  118. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +1 -1
  119. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +1 -1
  120. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +1 -1
  121. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +1 -1
  122. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer.pyi +3 -3
  123. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer_objects.pyi +3 -3
  124. metaflow-stubs/plugins/azure/__init__.pyi +1 -1
  125. metaflow-stubs/plugins/azure/azure_credential.pyi +1 -1
  126. metaflow-stubs/plugins/azure/azure_exceptions.pyi +1 -1
  127. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +2 -2
  128. metaflow-stubs/plugins/azure/azure_utils.pyi +1 -1
  129. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +1 -1
  130. metaflow-stubs/plugins/azure/includefile_support.pyi +1 -1
  131. metaflow-stubs/plugins/cards/__init__.pyi +1 -1
  132. metaflow-stubs/plugins/cards/card_client.pyi +2 -2
  133. metaflow-stubs/plugins/cards/card_creator.pyi +1 -1
  134. metaflow-stubs/plugins/cards/card_datastore.pyi +1 -1
  135. metaflow-stubs/plugins/cards/card_decorator.pyi +1 -1
  136. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +1 -1
  137. metaflow-stubs/plugins/cards/card_modules/basic.pyi +1 -1
  138. metaflow-stubs/plugins/cards/card_modules/card.pyi +1 -1
  139. metaflow-stubs/plugins/cards/card_modules/components.pyi +3 -3
  140. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +1 -1
  141. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +1 -1
  142. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +1 -1
  143. metaflow-stubs/plugins/cards/card_resolver.pyi +1 -1
  144. metaflow-stubs/plugins/cards/component_serializer.pyi +1 -1
  145. metaflow-stubs/plugins/cards/exception.pyi +1 -1
  146. metaflow-stubs/plugins/catch_decorator.pyi +1 -1
  147. metaflow-stubs/plugins/datatools/__init__.pyi +1 -1
  148. metaflow-stubs/plugins/datatools/local.pyi +1 -1
  149. metaflow-stubs/plugins/datatools/s3/__init__.pyi +1 -1
  150. metaflow-stubs/plugins/datatools/s3/s3.pyi +3 -3
  151. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +1 -1
  152. metaflow-stubs/plugins/datatools/s3/s3util.pyi +1 -1
  153. metaflow-stubs/plugins/debug_logger.pyi +1 -1
  154. metaflow-stubs/plugins/debug_monitor.pyi +1 -1
  155. metaflow-stubs/plugins/environment_decorator.pyi +1 -1
  156. metaflow-stubs/plugins/events_decorator.pyi +1 -1
  157. metaflow-stubs/plugins/frameworks/__init__.pyi +1 -1
  158. metaflow-stubs/plugins/frameworks/pytorch.pyi +1 -1
  159. metaflow-stubs/plugins/gcp/__init__.pyi +1 -1
  160. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +2 -2
  161. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +1 -1
  162. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +1 -1
  163. metaflow-stubs/plugins/gcp/gs_utils.pyi +1 -1
  164. metaflow-stubs/plugins/gcp/includefile_support.pyi +1 -1
  165. metaflow-stubs/plugins/kubernetes/__init__.pyi +1 -1
  166. metaflow-stubs/plugins/kubernetes/kube_utils.pyi +2 -2
  167. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +1 -1
  168. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +1 -1
  169. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +1 -1
  170. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +1 -1
  171. metaflow-stubs/plugins/kubernetes/spot_monitor_sidecar.pyi +1 -1
  172. metaflow-stubs/plugins/ollama/__init__.pyi +12 -3
  173. metaflow-stubs/plugins/parallel_decorator.pyi +1 -1
  174. metaflow-stubs/plugins/perimeters.pyi +1 -1
  175. metaflow-stubs/plugins/project_decorator.pyi +1 -1
  176. metaflow-stubs/plugins/pypi/__init__.pyi +2 -2
  177. metaflow-stubs/plugins/pypi/conda_decorator.pyi +1 -1
  178. metaflow-stubs/plugins/pypi/conda_environment.pyi +3 -3
  179. metaflow-stubs/plugins/pypi/parsers.pyi +1 -1
  180. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +1 -1
  181. metaflow-stubs/plugins/pypi/pypi_environment.pyi +1 -1
  182. metaflow-stubs/plugins/pypi/utils.pyi +1 -1
  183. metaflow-stubs/plugins/resources_decorator.pyi +1 -1
  184. metaflow-stubs/plugins/retry_decorator.pyi +1 -1
  185. metaflow-stubs/plugins/secrets/__init__.pyi +2 -2
  186. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +2 -2
  187. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +1 -1
  188. metaflow-stubs/plugins/snowflake/__init__.pyi +1 -1
  189. metaflow-stubs/plugins/storage_executor.pyi +1 -1
  190. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +1 -1
  191. metaflow-stubs/plugins/timeout_decorator.pyi +1 -1
  192. metaflow-stubs/plugins/torchtune/__init__.pyi +1 -1
  193. metaflow-stubs/plugins/uv/__init__.pyi +1 -1
  194. metaflow-stubs/plugins/uv/uv_environment.pyi +2 -2
  195. metaflow-stubs/profilers/__init__.pyi +1 -1
  196. metaflow-stubs/pylint_wrapper.pyi +1 -1
  197. metaflow-stubs/runner/__init__.pyi +1 -1
  198. metaflow-stubs/runner/deployer.pyi +4 -4
  199. metaflow-stubs/runner/deployer_impl.pyi +2 -2
  200. metaflow-stubs/runner/metaflow_runner.pyi +3 -3
  201. metaflow-stubs/runner/nbdeploy.pyi +1 -1
  202. metaflow-stubs/runner/nbrun.pyi +1 -1
  203. metaflow-stubs/runner/subprocess_manager.pyi +1 -1
  204. metaflow-stubs/runner/utils.pyi +3 -3
  205. metaflow-stubs/system/__init__.pyi +1 -1
  206. metaflow-stubs/system/system_logger.pyi +2 -2
  207. metaflow-stubs/system/system_monitor.pyi +1 -1
  208. metaflow-stubs/tagging_util.pyi +1 -1
  209. metaflow-stubs/tuple_util.pyi +1 -1
  210. metaflow-stubs/user_configs/__init__.pyi +1 -1
  211. metaflow-stubs/user_configs/config_decorators.pyi +4 -4
  212. metaflow-stubs/user_configs/config_options.pyi +3 -3
  213. metaflow-stubs/user_configs/config_parameters.pyi +6 -6
  214. {ob_metaflow_stubs-6.0.3.175rc1.dist-info → ob_metaflow_stubs-6.0.3.176rc1.dist-info}/METADATA +1 -1
  215. ob_metaflow_stubs-6.0.3.176rc1.dist-info/RECORD +218 -0
  216. ob_metaflow_stubs-6.0.3.175rc1.dist-info/RECORD +0 -216
  217. {ob_metaflow_stubs-6.0.3.175rc1.dist-info → ob_metaflow_stubs-6.0.3.176rc1.dist-info}/WHEEL +0 -0
  218. {ob_metaflow_stubs-6.0.3.175rc1.dist-info → ob_metaflow_stubs-6.0.3.176rc1.dist-info}/top_level.txt +0 -0
@@ -1,15 +1,15 @@
1
1
  ######################################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
3
  # MF version: 2.15.14.1+obcheckpoint(0.2.1);ob(v1) #
4
- # Generated on 2025-05-31T01:09:57.775269 #
4
+ # Generated on 2025-06-03T02:57:26.607647 #
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 tuple_util as tuple_util
39
38
  from . import cards as cards
39
+ from . import tuple_util as tuple_util
40
40
  from . import events as events
41
41
  from . import metaflow_git as metaflow_git
42
42
  from . import runner as runner
@@ -44,9 +44,9 @@ from . import plugins as plugins
44
44
  from .mf_extensions.outerbounds.toplevel.global_aliases_for_metaflow_package import S3 as S3
45
45
  from . import includefile as includefile
46
46
  from .includefile import IncludeFile as IncludeFile
47
+ from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
47
48
  from .plugins.pypi.parsers import conda_environment_yml_parser as conda_environment_yml_parser
48
49
  from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
49
- from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
50
50
  from . import client as client
51
51
  from .client.core import namespace as namespace
52
52
  from .client.core import get_namespace as get_namespace
@@ -154,82 +154,216 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
154
154
  """
155
155
  ...
156
156
 
157
+ def nvidia(*, gpu: int, gpu_type: str, queue_timeout: int) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
158
+ """
159
+ Specifies that this step should execute on DGX cloud.
160
+
161
+
162
+ Parameters
163
+ ----------
164
+ gpu : int
165
+ Number of GPUs to use.
166
+ gpu_type : str
167
+ Type of Nvidia GPU to use.
168
+ queue_timeout : int
169
+ Time to keep the job in NVCF's queue.
170
+ """
171
+ ...
172
+
157
173
  @typing.overload
158
- 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]]]:
174
+ 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]]]:
159
175
  """
160
- Specifies the resources needed when executing this step.
176
+ Enables checkpointing for a step.
161
177
 
162
- Use `@resources` to specify the resource requirements
163
- independently of the specific compute layer (`@batch`, `@kubernetes`).
164
178
 
165
- You can choose the compute layer on the command line by executing e.g.
166
- ```
167
- python myflow.py run --with batch
168
- ```
169
- or
170
- ```
171
- python myflow.py run --with kubernetes
172
- ```
173
- which executes the flow on the desired system using the
174
- requirements specified in `@resources`.
179
+
180
+ Parameters
181
+ ----------
182
+ load_policy : str, default: "fresh"
183
+ The policy for loading the checkpoint. The following policies are supported:
184
+ - "eager": Loads the the latest available checkpoint within the namespace.
185
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
186
+ will be loaded at the start of the task.
187
+ - "none": Do not load any checkpoint
188
+ - "fresh": Loads the lastest checkpoint created within the running Task.
189
+ This mode helps loading checkpoints across various retry attempts of the same task.
190
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
191
+ created within the task will be loaded when the task is retries execution on failure.
192
+
193
+ temp_dir_root : str, default: None
194
+ The root directory under which `current.checkpoint.directory` will be created.
195
+ """
196
+ ...
197
+
198
+ @typing.overload
199
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
200
+ ...
201
+
202
+ @typing.overload
203
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
204
+ ...
205
+
206
+ 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):
207
+ """
208
+ Enables checkpointing for a step.
209
+
175
210
 
176
211
 
177
212
  Parameters
178
213
  ----------
179
- cpu : int, default 1
180
- Number of CPUs required for this step.
181
- gpu : int, optional, default None
182
- Number of GPUs required for this step.
183
- disk : int, optional, default None
184
- Disk size (in MB) required for this step. Only applies on Kubernetes.
185
- memory : int, default 4096
186
- Memory size (in MB) required for this step.
187
- shared_memory : int, optional, default None
188
- The value for the size (in MiB) of the /dev/shm volume for this step.
189
- This parameter maps to the `--shm-size` option in Docker.
214
+ load_policy : str, default: "fresh"
215
+ The policy for loading the checkpoint. The following policies are supported:
216
+ - "eager": Loads the the latest available checkpoint within the namespace.
217
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
218
+ will be loaded at the start of the task.
219
+ - "none": Do not load any checkpoint
220
+ - "fresh": Loads the lastest checkpoint created within the running Task.
221
+ This mode helps loading checkpoints across various retry attempts of the same task.
222
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
223
+ created within the task will be loaded when the task is retries execution on failure.
224
+
225
+ temp_dir_root : str, default: None
226
+ The root directory under which `current.checkpoint.directory` will be created.
190
227
  """
191
228
  ...
192
229
 
193
230
  @typing.overload
194
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
231
+ def retry(*, times: int = 3, minutes_between_retries: int = 2) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
232
+ """
233
+ Specifies the number of times the task corresponding
234
+ to a step needs to be retried.
235
+
236
+ This decorator is useful for handling transient errors, such as networking issues.
237
+ If your task contains operations that can't be retried safely, e.g. database updates,
238
+ it is advisable to annotate it with `@retry(times=0)`.
239
+
240
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
241
+ decorator will execute a no-op task after all retries have been exhausted,
242
+ ensuring that the flow execution can continue.
243
+
244
+
245
+ Parameters
246
+ ----------
247
+ times : int, default 3
248
+ Number of times to retry this task.
249
+ minutes_between_retries : int, default 2
250
+ Number of minutes between retries.
251
+ """
195
252
  ...
196
253
 
197
254
  @typing.overload
198
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
255
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
199
256
  ...
200
257
 
201
- 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):
258
+ @typing.overload
259
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
260
+ ...
261
+
262
+ def retry(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, times: int = 3, minutes_between_retries: int = 2):
202
263
  """
203
- Specifies the resources needed when executing this step.
264
+ Specifies the number of times the task corresponding
265
+ to a step needs to be retried.
204
266
 
205
- Use `@resources` to specify the resource requirements
206
- independently of the specific compute layer (`@batch`, `@kubernetes`).
267
+ This decorator is useful for handling transient errors, such as networking issues.
268
+ If your task contains operations that can't be retried safely, e.g. database updates,
269
+ it is advisable to annotate it with `@retry(times=0)`.
207
270
 
208
- You can choose the compute layer on the command line by executing e.g.
209
- ```
210
- python myflow.py run --with batch
211
- ```
212
- or
213
- ```
214
- python myflow.py run --with kubernetes
215
- ```
216
- which executes the flow on the desired system using the
217
- requirements specified in `@resources`.
271
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
272
+ decorator will execute a no-op task after all retries have been exhausted,
273
+ ensuring that the flow execution can continue.
274
+
275
+
276
+ Parameters
277
+ ----------
278
+ times : int, default 3
279
+ Number of times to retry this task.
280
+ minutes_between_retries : int, default 2
281
+ Number of minutes between retries.
282
+ """
283
+ ...
284
+
285
+ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: typing.Optional[str] = None, image_pull_policy: str = 'KUBERNETES_IMAGE_PULL_POLICY', service_account: str = 'METAFLOW_KUBERNETES_SERVICE_ACCOUNT', secrets: typing.Optional[typing.List[str]] = None, node_selector: typing.Union[typing.Dict[str, str], str, None] = None, namespace: str = 'METAFLOW_KUBERNETES_NAMESPACE', gpu: typing.Optional[int] = None, gpu_vendor: str = 'KUBERNETES_GPU_VENDOR', tolerations: typing.List[str] = [], labels: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_LABELS', annotations: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_ANNOTATIONS', use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = '/metaflow_temp', persistent_volume_claims: typing.Optional[typing.Dict[str, str]] = None, shared_memory: typing.Optional[int] = None, port: typing.Optional[int] = None, compute_pool: typing.Optional[str] = None, hostname_resolution_timeout: int = 600, qos: str = 'Burstable', security_context: typing.Optional[typing.Dict[str, typing.Any]] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
286
+ """
287
+ Specifies that this step should execute on Kubernetes.
218
288
 
219
289
 
220
290
  Parameters
221
291
  ----------
222
292
  cpu : int, default 1
223
- Number of CPUs required for this step.
224
- gpu : int, optional, default None
225
- Number of GPUs required for this step.
226
- disk : int, optional, default None
227
- Disk size (in MB) required for this step. Only applies on Kubernetes.
293
+ Number of CPUs required for this step. If `@resources` is
294
+ also present, the maximum value from all decorators is used.
228
295
  memory : int, default 4096
229
- Memory size (in MB) required for this step.
230
- shared_memory : int, optional, default None
231
- The value for the size (in MiB) of the /dev/shm volume for this step.
232
- This parameter maps to the `--shm-size` option in Docker.
296
+ Memory size (in MB) required for this step. If
297
+ `@resources` is also present, the maximum value from all decorators is
298
+ used.
299
+ disk : int, default 10240
300
+ Disk size (in MB) required for this step. If
301
+ `@resources` is also present, the maximum value from all decorators is
302
+ used.
303
+ image : str, optional, default None
304
+ Docker image to use when launching on Kubernetes. If not specified, and
305
+ METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
306
+ not, a default Docker image mapping to the current version of Python is used.
307
+ image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
308
+ If given, the imagePullPolicy to be applied to the Docker image of the step.
309
+ service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
310
+ Kubernetes service account to use when launching pod in Kubernetes.
311
+ secrets : List[str], optional, default None
312
+ Kubernetes secrets to use when launching pod in Kubernetes. These
313
+ secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
314
+ in Metaflow configuration.
315
+ node_selector: Union[Dict[str,str], str], optional, default None
316
+ Kubernetes node selector(s) to apply to the pod running the task.
317
+ Can be passed in as a comma separated string of values e.g.
318
+ 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
319
+ {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
320
+ namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
321
+ Kubernetes namespace to use when launching pod in Kubernetes.
322
+ gpu : int, optional, default None
323
+ Number of GPUs required for this step. A value of zero implies that
324
+ the scheduled node should not have GPUs.
325
+ gpu_vendor : str, default KUBERNETES_GPU_VENDOR
326
+ The vendor of the GPUs to be used for this step.
327
+ tolerations : List[str], default []
328
+ The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
329
+ Kubernetes tolerations to use when launching pod in Kubernetes.
330
+ labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
331
+ Kubernetes labels to use when launching pod in Kubernetes.
332
+ annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
333
+ Kubernetes annotations to use when launching pod in Kubernetes.
334
+ use_tmpfs : bool, default False
335
+ This enables an explicit tmpfs mount for this step.
336
+ tmpfs_tempdir : bool, default True
337
+ sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
338
+ tmpfs_size : int, optional, default: None
339
+ The value for the size (in MiB) of the tmpfs mount for this step.
340
+ This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
341
+ memory allocated for this step.
342
+ tmpfs_path : str, optional, default /metaflow_temp
343
+ Path to tmpfs mount for this step.
344
+ persistent_volume_claims : Dict[str, str], optional, default None
345
+ A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
346
+ volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
347
+ shared_memory: int, optional
348
+ Shared memory size (in MiB) required for this step
349
+ port: int, optional
350
+ Port number to specify in the Kubernetes job object
351
+ compute_pool : str, optional, default None
352
+ Compute pool to be used for for this step.
353
+ If not specified, any accessible compute pool within the perimeter is used.
354
+ hostname_resolution_timeout: int, default 10 * 60
355
+ Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
356
+ Only applicable when @parallel is used.
357
+ qos: str, default: Burstable
358
+ Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
359
+
360
+ security_context: Dict[str, Any], optional, default None
361
+ Container security context. Applies to the task container. Allows the following keys:
362
+ - privileged: bool, optional, default None
363
+ - allow_privilege_escalation: bool, optional, default None
364
+ - run_as_user: int, optional, default None
365
+ - run_as_group: int, optional, default None
366
+ - run_as_non_root: bool, optional, default None
233
367
  """
234
368
  ...
235
369
 
@@ -284,64 +418,32 @@ def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
284
418
  """
285
419
  ...
286
420
 
287
- @typing.overload
288
- 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]]]:
421
+ 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]]]:
289
422
  """
290
- Enables loading / saving of models within a step.
291
-
423
+ Decorator that helps cache, version and store models/datasets from huggingface hub.
292
424
 
293
425
 
294
426
  Parameters
295
427
  ----------
296
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
297
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
298
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
299
- - `current.checkpoint`
300
- - `current.model`
301
- - `current.huggingface_hub`
302
-
303
- 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
304
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
305
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
306
-
307
- temp_dir_root : str, default: None
308
- The root directory under which `current.model.loaded` will store loaded models
309
- """
310
- ...
311
-
312
- @typing.overload
313
- def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
314
- ...
315
-
316
- @typing.overload
317
- def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
318
- ...
319
-
320
- 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):
321
- """
322
- Enables loading / saving of models within a step.
323
-
428
+ temp_dir_root : str, optional
429
+ The root directory that will hold the temporary directory where objects will be downloaded.
324
430
 
431
+ load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
432
+ The list of repos (models/datasets) to load.
325
433
 
326
- Parameters
327
- ----------
328
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
329
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
330
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
331
- - `current.checkpoint`
332
- - `current.model`
333
- - `current.huggingface_hub`
434
+ Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
334
435
 
335
- 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
336
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
337
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
436
+ - If repo (model/dataset) is not found in the datastore:
437
+ - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
438
+ - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
439
+ - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
338
440
 
339
- temp_dir_root : str, default: None
340
- The root directory under which `current.model.loaded` will store loaded models
441
+ - If repo is found in the datastore:
442
+ - Loads it directly from datastore to local path (can be temporary directory or specified path)
341
443
  """
342
444
  ...
343
445
 
344
- def ollama(*, models: list, backend: str, force_pull: bool, skip_push_check: bool, debug: bool) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
446
+ def ollama(*, models: list, backend: str, force_pull: bool, cache_update_policy: str, force_cache_update: bool, debug: bool, circuit_breaker_config: dict, timeout_config: dict) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
345
447
  """
346
448
  This decorator is used to run Ollama APIs as Metaflow task sidecars.
347
449
 
@@ -371,348 +473,255 @@ def ollama(*, models: list, backend: str, force_pull: bool, skip_push_check: boo
371
473
  Determines where and how to run the Ollama process.
372
474
  force_pull: bool
373
475
  Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
374
- skip_push_check: bool
375
- Whether to skip the check that populates/overwrites remote cache on terminating an ollama model.
476
+ cache_update_policy: str
477
+ Cache update policy: "auto", "force", or "never".
478
+ force_cache_update: bool
479
+ Simple override for "force" cache update policy.
376
480
  debug: bool
377
481
  Whether to turn on verbose debugging logs.
482
+ circuit_breaker_config: dict
483
+ Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
484
+ timeout_config: dict
485
+ Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
378
486
  """
379
487
  ...
380
488
 
381
- def app_deploy(*, app_port: int, app_name: str) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
489
+ @typing.overload
490
+ 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]]]:
382
491
  """
383
- Specifies that this step is used to deploy an instance of the app.
384
- Requires that self.app_name, self.app_port, self.entrypoint and self.deployDir is set.
492
+ Specifies a timeout for your step.
493
+
494
+ This decorator is useful if this step may hang indefinitely.
495
+
496
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
497
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
498
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
499
+
500
+ Note that all the values specified in parameters are added together so if you specify
501
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
385
502
 
386
503
 
387
504
  Parameters
388
505
  ----------
389
- app_port : int
390
- Number of GPUs to use.
391
- app_name : str
392
- Name of the app to deploy.
393
- """
394
- ...
395
-
396
- @typing.overload
397
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
398
- """
399
- Decorator prototype for all step decorators. This function gets specialized
400
- and imported for all decorators types by _import_plugin_decorators().
401
- """
402
- ...
403
-
404
- @typing.overload
405
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
406
- ...
407
-
408
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
409
- """
410
- Decorator prototype for all step decorators. This function gets specialized
411
- and imported for all decorators types by _import_plugin_decorators().
506
+ seconds : int, default 0
507
+ Number of seconds to wait prior to timing out.
508
+ minutes : int, default 0
509
+ Number of minutes to wait prior to timing out.
510
+ hours : int, default 0
511
+ Number of hours to wait prior to timing out.
412
512
  """
413
513
  ...
414
514
 
415
515
  @typing.overload
416
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
417
- """
418
- Internal decorator to support Fast bakery
419
- """
516
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
420
517
  ...
421
518
 
422
519
  @typing.overload
423
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
424
- ...
425
-
426
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
427
- """
428
- Internal decorator to support Fast bakery
429
- """
430
- ...
431
-
432
- def nvct(*, gpu: int, gpu_type: str) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
433
- """
434
- Specifies that this step should execute on DGX cloud.
435
-
436
-
437
- Parameters
438
- ----------
439
- gpu : int
440
- Number of GPUs to use.
441
- gpu_type : str
442
- Type of Nvidia GPU to use.
443
- """
520
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
444
521
  ...
445
522
 
446
- def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: typing.Optional[str] = None, image_pull_policy: str = 'KUBERNETES_IMAGE_PULL_POLICY', service_account: str = 'METAFLOW_KUBERNETES_SERVICE_ACCOUNT', secrets: typing.Optional[typing.List[str]] = None, node_selector: typing.Union[typing.Dict[str, str], str, None] = None, namespace: str = 'METAFLOW_KUBERNETES_NAMESPACE', gpu: typing.Optional[int] = None, gpu_vendor: str = 'KUBERNETES_GPU_VENDOR', tolerations: typing.List[str] = [], labels: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_LABELS', annotations: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_ANNOTATIONS', use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = '/metaflow_temp', persistent_volume_claims: typing.Optional[typing.Dict[str, str]] = None, shared_memory: typing.Optional[int] = None, port: typing.Optional[int] = None, compute_pool: typing.Optional[str] = None, hostname_resolution_timeout: int = 600, qos: str = 'Burstable', security_context: typing.Optional[typing.Dict[str, typing.Any]] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
523
+ 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):
447
524
  """
448
- Specifies that this step should execute on Kubernetes.
525
+ Specifies a timeout for your step.
449
526
 
527
+ This decorator is useful if this step may hang indefinitely.
450
528
 
451
- Parameters
452
- ----------
453
- cpu : int, default 1
454
- Number of CPUs required for this step. If `@resources` is
455
- also present, the maximum value from all decorators is used.
456
- memory : int, default 4096
457
- Memory size (in MB) required for this step. If
458
- `@resources` is also present, the maximum value from all decorators is
459
- used.
460
- disk : int, default 10240
461
- Disk size (in MB) required for this step. If
462
- `@resources` is also present, the maximum value from all decorators is
463
- used.
464
- image : str, optional, default None
465
- Docker image to use when launching on Kubernetes. If not specified, and
466
- METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
467
- not, a default Docker image mapping to the current version of Python is used.
468
- image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
469
- If given, the imagePullPolicy to be applied to the Docker image of the step.
470
- service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
471
- Kubernetes service account to use when launching pod in Kubernetes.
472
- secrets : List[str], optional, default None
473
- Kubernetes secrets to use when launching pod in Kubernetes. These
474
- secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
475
- in Metaflow configuration.
476
- node_selector: Union[Dict[str,str], str], optional, default None
477
- Kubernetes node selector(s) to apply to the pod running the task.
478
- Can be passed in as a comma separated string of values e.g.
479
- 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
480
- {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
481
- namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
482
- Kubernetes namespace to use when launching pod in Kubernetes.
483
- gpu : int, optional, default None
484
- Number of GPUs required for this step. A value of zero implies that
485
- the scheduled node should not have GPUs.
486
- gpu_vendor : str, default KUBERNETES_GPU_VENDOR
487
- The vendor of the GPUs to be used for this step.
488
- tolerations : List[str], default []
489
- The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
490
- Kubernetes tolerations to use when launching pod in Kubernetes.
491
- labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
492
- Kubernetes labels to use when launching pod in Kubernetes.
493
- annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
494
- Kubernetes annotations to use when launching pod in Kubernetes.
495
- use_tmpfs : bool, default False
496
- This enables an explicit tmpfs mount for this step.
497
- tmpfs_tempdir : bool, default True
498
- sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
499
- tmpfs_size : int, optional, default: None
500
- The value for the size (in MiB) of the tmpfs mount for this step.
501
- This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
502
- memory allocated for this step.
503
- tmpfs_path : str, optional, default /metaflow_temp
504
- Path to tmpfs mount for this step.
505
- persistent_volume_claims : Dict[str, str], optional, default None
506
- A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
507
- volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
508
- shared_memory: int, optional
509
- Shared memory size (in MiB) required for this step
510
- port: int, optional
511
- Port number to specify in the Kubernetes job object
512
- compute_pool : str, optional, default None
513
- Compute pool to be used for for this step.
514
- If not specified, any accessible compute pool within the perimeter is used.
515
- hostname_resolution_timeout: int, default 10 * 60
516
- Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
517
- Only applicable when @parallel is used.
518
- qos: str, default: Burstable
519
- Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
529
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
530
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
531
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
520
532
 
521
- security_context: Dict[str, Any], optional, default None
522
- Container security context. Applies to the task container. Allows the following keys:
523
- - privileged: bool, optional, default None
524
- - allow_privilege_escalation: bool, optional, default None
525
- - run_as_user: int, optional, default None
526
- - run_as_group: int, optional, default None
527
- - run_as_non_root: bool, optional, default None
528
- """
529
- ...
530
-
531
- 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]]]:
532
- """
533
- Specifies that this step should execute on DGX cloud.
533
+ Note that all the values specified in parameters are added together so if you specify
534
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
534
535
 
535
536
 
536
537
  Parameters
537
538
  ----------
538
- gpu : int
539
- Number of GPUs to use.
540
- gpu_type : str
541
- Type of Nvidia GPU to use.
542
- queue_timeout : int
543
- Time to keep the job in NVCF's queue.
539
+ seconds : int, default 0
540
+ Number of seconds to wait prior to timing out.
541
+ minutes : int, default 0
542
+ Number of minutes to wait prior to timing out.
543
+ hours : int, default 0
544
+ Number of hours to wait prior to timing out.
544
545
  """
545
546
  ...
546
547
 
547
548
  @typing.overload
548
- 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]]]:
549
+ 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]]]:
549
550
  """
550
- Specifies environment variables to be set prior to the execution of a step.
551
+ Specifies the resources needed when executing this step.
552
+
553
+ Use `@resources` to specify the resource requirements
554
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
555
+
556
+ You can choose the compute layer on the command line by executing e.g.
557
+ ```
558
+ python myflow.py run --with batch
559
+ ```
560
+ or
561
+ ```
562
+ python myflow.py run --with kubernetes
563
+ ```
564
+ which executes the flow on the desired system using the
565
+ requirements specified in `@resources`.
551
566
 
552
567
 
553
568
  Parameters
554
569
  ----------
555
- vars : Dict[str, str], default {}
556
- Dictionary of environment variables to set.
570
+ cpu : int, default 1
571
+ Number of CPUs required for this step.
572
+ gpu : int, optional, default None
573
+ Number of GPUs required for this step.
574
+ disk : int, optional, default None
575
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
576
+ memory : int, default 4096
577
+ Memory size (in MB) required for this step.
578
+ shared_memory : int, optional, default None
579
+ The value for the size (in MiB) of the /dev/shm volume for this step.
580
+ This parameter maps to the `--shm-size` option in Docker.
557
581
  """
558
582
  ...
559
583
 
560
584
  @typing.overload
561
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
585
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
562
586
  ...
563
587
 
564
588
  @typing.overload
565
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
589
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
566
590
  ...
567
591
 
568
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
592
+ 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):
569
593
  """
570
- Specifies environment variables to be set prior to the execution of a step.
594
+ Specifies the resources needed when executing this step.
595
+
596
+ Use `@resources` to specify the resource requirements
597
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
598
+
599
+ You can choose the compute layer on the command line by executing e.g.
600
+ ```
601
+ python myflow.py run --with batch
602
+ ```
603
+ or
604
+ ```
605
+ python myflow.py run --with kubernetes
606
+ ```
607
+ which executes the flow on the desired system using the
608
+ requirements specified in `@resources`.
571
609
 
572
610
 
573
611
  Parameters
574
612
  ----------
575
- vars : Dict[str, str], default {}
576
- Dictionary of environment variables to set.
613
+ cpu : int, default 1
614
+ Number of CPUs required for this step.
615
+ gpu : int, optional, default None
616
+ Number of GPUs required for this step.
617
+ disk : int, optional, default None
618
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
619
+ memory : int, default 4096
620
+ Memory size (in MB) required for this step.
621
+ shared_memory : int, optional, default None
622
+ The value for the size (in MiB) of the /dev/shm volume for this step.
623
+ This parameter maps to the `--shm-size` option in Docker.
577
624
  """
578
625
  ...
579
626
 
580
- 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]]]:
627
+ @typing.overload
628
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
581
629
  """
582
- Decorator that helps cache, version and store models/datasets from huggingface hub.
583
-
584
-
585
- Parameters
586
- ----------
587
- temp_dir_root : str, optional
588
- The root directory that will hold the temporary directory where objects will be downloaded.
589
-
590
- load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
591
- The list of repos (models/datasets) to load.
592
-
593
- Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
594
-
595
- - If repo (model/dataset) is not found in the datastore:
596
- - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
597
- - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
598
- - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
599
-
600
- - If repo is found in the datastore:
601
- - Loads it directly from datastore to local path (can be temporary directory or specified path)
630
+ Internal decorator to support Fast bakery
602
631
  """
603
632
  ...
604
633
 
605
634
  @typing.overload
606
- def retry(*, times: int = 3, minutes_between_retries: int = 2) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
635
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
636
+ ...
637
+
638
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
607
639
  """
608
- Specifies the number of times the task corresponding
609
- to a step needs to be retried.
610
-
611
- This decorator is useful for handling transient errors, such as networking issues.
612
- If your task contains operations that can't be retried safely, e.g. database updates,
613
- it is advisable to annotate it with `@retry(times=0)`.
614
-
615
- This can be used in conjunction with the `@catch` decorator. The `@catch`
616
- decorator will execute a no-op task after all retries have been exhausted,
617
- ensuring that the flow execution can continue.
640
+ Internal decorator to support Fast bakery
641
+ """
642
+ ...
643
+
644
+ @typing.overload
645
+ 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]]]:
646
+ """
647
+ Specifies environment variables to be set prior to the execution of a step.
618
648
 
619
649
 
620
650
  Parameters
621
651
  ----------
622
- times : int, default 3
623
- Number of times to retry this task.
624
- minutes_between_retries : int, default 2
625
- Number of minutes between retries.
652
+ vars : Dict[str, str], default {}
653
+ Dictionary of environment variables to set.
626
654
  """
627
655
  ...
628
656
 
629
657
  @typing.overload
630
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
658
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
631
659
  ...
632
660
 
633
661
  @typing.overload
634
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
662
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
635
663
  ...
636
664
 
637
- def retry(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, times: int = 3, minutes_between_retries: int = 2):
665
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
638
666
  """
639
- Specifies the number of times the task corresponding
640
- to a step needs to be retried.
641
-
642
- This decorator is useful for handling transient errors, such as networking issues.
643
- If your task contains operations that can't be retried safely, e.g. database updates,
644
- it is advisable to annotate it with `@retry(times=0)`.
645
-
646
- This can be used in conjunction with the `@catch` decorator. The `@catch`
647
- decorator will execute a no-op task after all retries have been exhausted,
648
- ensuring that the flow execution can continue.
667
+ Specifies environment variables to be set prior to the execution of a step.
649
668
 
650
669
 
651
670
  Parameters
652
671
  ----------
653
- times : int, default 3
654
- Number of times to retry this task.
655
- minutes_between_retries : int, default 2
656
- Number of minutes between retries.
672
+ vars : Dict[str, str], default {}
673
+ Dictionary of environment variables to set.
657
674
  """
658
675
  ...
659
676
 
660
677
  @typing.overload
661
- 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]]]:
678
+ 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]]]:
662
679
  """
663
- Specifies a timeout for your step.
664
-
665
- This decorator is useful if this step may hang indefinitely.
666
-
667
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
668
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
669
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
680
+ Specifies the PyPI packages for the step.
670
681
 
671
- Note that all the values specified in parameters are added together so if you specify
672
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
682
+ Information in this decorator will augment any
683
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
684
+ you can use `@pypi_base` to set packages required by all
685
+ steps and use `@pypi` to specify step-specific overrides.
673
686
 
674
687
 
675
688
  Parameters
676
689
  ----------
677
- seconds : int, default 0
678
- Number of seconds to wait prior to timing out.
679
- minutes : int, default 0
680
- Number of minutes to wait prior to timing out.
681
- hours : int, default 0
682
- Number of hours to wait prior to timing out.
690
+ packages : Dict[str, str], default: {}
691
+ Packages to use for this step. The key is the name of the package
692
+ and the value is the version to use.
693
+ python : str, optional, default: None
694
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
695
+ that the version used will correspond to the version of the Python interpreter used to start the run.
683
696
  """
684
697
  ...
685
698
 
686
699
  @typing.overload
687
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
700
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
688
701
  ...
689
702
 
690
703
  @typing.overload
691
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
704
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
692
705
  ...
693
706
 
694
- 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):
707
+ 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):
695
708
  """
696
- Specifies a timeout for your step.
697
-
698
- This decorator is useful if this step may hang indefinitely.
699
-
700
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
701
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
702
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
709
+ Specifies the PyPI packages for the step.
703
710
 
704
- Note that all the values specified in parameters are added together so if you specify
705
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
711
+ Information in this decorator will augment any
712
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
713
+ you can use `@pypi_base` to set packages required by all
714
+ steps and use `@pypi` to specify step-specific overrides.
706
715
 
707
716
 
708
717
  Parameters
709
718
  ----------
710
- seconds : int, default 0
711
- Number of seconds to wait prior to timing out.
712
- minutes : int, default 0
713
- Number of minutes to wait prior to timing out.
714
- hours : int, default 0
715
- Number of hours to wait prior to timing out.
719
+ packages : Dict[str, str], default: {}
720
+ Packages to use for this step. The key is the name of the package
721
+ and the value is the version to use.
722
+ python : str, optional, default: None
723
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
724
+ that the version used will correspond to the version of the Python interpreter used to start the run.
716
725
  """
717
726
  ...
718
727
 
@@ -766,118 +775,107 @@ def card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typ
766
775
  ...
767
776
 
768
777
  @typing.overload
769
- 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]]]:
778
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
770
779
  """
771
- Enables checkpointing for a step.
772
-
773
-
774
-
775
- Parameters
776
- ----------
777
- load_policy : str, default: "fresh"
778
- The policy for loading the checkpoint. The following policies are supported:
779
- - "eager": Loads the the latest available checkpoint within the namespace.
780
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
781
- will be loaded at the start of the task.
782
- - "none": Do not load any checkpoint
783
- - "fresh": Loads the lastest checkpoint created within the running Task.
784
- This mode helps loading checkpoints across various retry attempts of the same task.
785
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
786
- created within the task will be loaded when the task is retries execution on failure.
787
-
788
- temp_dir_root : str, default: None
789
- The root directory under which `current.checkpoint.directory` will be created.
780
+ Decorator prototype for all step decorators. This function gets specialized
781
+ and imported for all decorators types by _import_plugin_decorators().
790
782
  """
791
783
  ...
792
784
 
793
785
  @typing.overload
794
- def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
786
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
795
787
  ...
796
788
 
797
- @typing.overload
798
- def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
789
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
790
+ """
791
+ Decorator prototype for all step decorators. This function gets specialized
792
+ and imported for all decorators types by _import_plugin_decorators().
793
+ """
799
794
  ...
800
795
 
801
- 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):
796
+ def nvct(*, gpu: int, gpu_type: str) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
802
797
  """
803
- Enables checkpointing for a step.
804
-
798
+ Specifies that this step should execute on DGX cloud.
805
799
 
806
800
 
807
801
  Parameters
808
802
  ----------
809
- load_policy : str, default: "fresh"
810
- The policy for loading the checkpoint. The following policies are supported:
811
- - "eager": Loads the the latest available checkpoint within the namespace.
812
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
813
- will be loaded at the start of the task.
814
- - "none": Do not load any checkpoint
815
- - "fresh": Loads the lastest checkpoint created within the running Task.
816
- This mode helps loading checkpoints across various retry attempts of the same task.
817
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
818
- created within the task will be loaded when the task is retries execution on failure.
803
+ gpu : int
804
+ Number of GPUs to use.
805
+ gpu_type : str
806
+ Type of Nvidia GPU to use.
807
+ """
808
+ ...
809
+
810
+ def app_deploy(*, app_port: int, app_name: str) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
811
+ """
812
+ Specifies that this step is used to deploy an instance of the app.
813
+ Requires that self.app_name, self.app_port, self.entrypoint and self.deployDir is set.
819
814
 
820
- temp_dir_root : str, default: None
821
- The root directory under which `current.checkpoint.directory` will be created.
815
+
816
+ Parameters
817
+ ----------
818
+ app_port : int
819
+ Number of GPUs to use.
820
+ app_name : str
821
+ Name of the app to deploy.
822
822
  """
823
823
  ...
824
824
 
825
825
  @typing.overload
826
- 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]]]:
826
+ 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]]]:
827
827
  """
828
- Specifies the Conda environment for the step.
828
+ Enables loading / saving of models within a step.
829
829
 
830
- Information in this decorator will augment any
831
- attributes set in the `@conda_base` flow-level decorator. Hence,
832
- you can use `@conda_base` to set packages required by all
833
- steps and use `@conda` to specify step-specific overrides.
834
830
 
835
831
 
836
832
  Parameters
837
833
  ----------
838
- packages : Dict[str, str], default {}
839
- Packages to use for this step. The key is the name of the package
840
- and the value is the version to use.
841
- libraries : Dict[str, str], default {}
842
- Supported for backward compatibility. When used with packages, packages will take precedence.
843
- python : str, optional, default None
844
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
845
- that the version used will correspond to the version of the Python interpreter used to start the run.
846
- disabled : bool, default False
847
- If set to True, disables @conda.
834
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
835
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
836
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
837
+ - `current.checkpoint`
838
+ - `current.model`
839
+ - `current.huggingface_hub`
840
+
841
+ 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
842
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
843
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
844
+
845
+ temp_dir_root : str, default: None
846
+ The root directory under which `current.model.loaded` will store loaded models
848
847
  """
849
848
  ...
850
849
 
851
850
  @typing.overload
852
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
851
+ def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
853
852
  ...
854
853
 
855
854
  @typing.overload
856
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
855
+ def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
857
856
  ...
858
857
 
859
- 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):
858
+ 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):
860
859
  """
861
- Specifies the Conda environment for the step.
860
+ Enables loading / saving of models within a step.
862
861
 
863
- Information in this decorator will augment any
864
- attributes set in the `@conda_base` flow-level decorator. Hence,
865
- you can use `@conda_base` to set packages required by all
866
- steps and use `@conda` to specify step-specific overrides.
867
862
 
868
863
 
869
864
  Parameters
870
865
  ----------
871
- packages : Dict[str, str], default {}
872
- Packages to use for this step. The key is the name of the package
873
- and the value is the version to use.
874
- libraries : Dict[str, str], default {}
875
- Supported for backward compatibility. When used with packages, packages will take precedence.
876
- python : str, optional, default None
877
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
878
- that the version used will correspond to the version of the Python interpreter used to start the run.
879
- disabled : bool, default False
880
- If set to True, disables @conda.
866
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
867
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
868
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
869
+ - `current.checkpoint`
870
+ - `current.model`
871
+ - `current.huggingface_hub`
872
+
873
+ 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
874
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
875
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
876
+
877
+ temp_dir_root : str, default: None
878
+ The root directory under which `current.model.loaded` will store loaded models
881
879
  """
882
880
  ...
883
881
 
@@ -917,341 +915,417 @@ def secrets(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None],
917
915
  ...
918
916
 
919
917
  @typing.overload
920
- 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]]]:
918
+ 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]]]:
919
+ """
920
+ Specifies the Conda environment for the step.
921
+
922
+ Information in this decorator will augment any
923
+ attributes set in the `@conda_base` flow-level decorator. Hence,
924
+ you can use `@conda_base` to set packages required by all
925
+ steps and use `@conda` to specify step-specific overrides.
926
+
927
+
928
+ Parameters
929
+ ----------
930
+ packages : Dict[str, str], default {}
931
+ Packages to use for this step. The key is the name of the package
932
+ and the value is the version to use.
933
+ libraries : Dict[str, str], default {}
934
+ Supported for backward compatibility. When used with packages, packages will take precedence.
935
+ python : str, optional, default None
936
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
937
+ that the version used will correspond to the version of the Python interpreter used to start the run.
938
+ disabled : bool, default False
939
+ If set to True, disables @conda.
940
+ """
941
+ ...
942
+
943
+ @typing.overload
944
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
945
+ ...
946
+
947
+ @typing.overload
948
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
949
+ ...
950
+
951
+ 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):
921
952
  """
922
- Specifies the PyPI packages for the step.
953
+ Specifies the Conda environment for the step.
923
954
 
924
955
  Information in this decorator will augment any
925
- attributes set in the `@pyi_base` flow-level decorator. Hence,
926
- you can use `@pypi_base` to set packages required by all
927
- steps and use `@pypi` to specify step-specific overrides.
956
+ attributes set in the `@conda_base` flow-level decorator. Hence,
957
+ you can use `@conda_base` to set packages required by all
958
+ steps and use `@conda` to specify step-specific overrides.
928
959
 
929
960
 
930
961
  Parameters
931
962
  ----------
932
- packages : Dict[str, str], default: {}
963
+ packages : Dict[str, str], default {}
933
964
  Packages to use for this step. The key is the name of the package
934
965
  and the value is the version to use.
935
- python : str, optional, default: None
966
+ libraries : Dict[str, str], default {}
967
+ Supported for backward compatibility. When used with packages, packages will take precedence.
968
+ python : str, optional, default None
936
969
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
937
970
  that the version used will correspond to the version of the Python interpreter used to start the run.
971
+ disabled : bool, default False
972
+ If set to True, disables @conda.
938
973
  """
939
974
  ...
940
975
 
941
- @typing.overload
942
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
976
+ def airflow_external_task_sensor(*, timeout: int, poke_interval: int, mode: str, exponential_backoff: bool, pool: str, soft_fail: bool, name: str, description: str, external_dag_id: str, external_task_ids: typing.List[str], allowed_states: typing.List[str], failed_states: typing.List[str], execution_delta: "datetime.timedelta", check_existence: bool) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
977
+ """
978
+ The `@airflow_external_task_sensor` decorator attaches a Airflow [ExternalTaskSensor](https://airflow.apache.org/docs/apache-airflow/stable/_api/airflow/sensors/external_task/index.html#airflow.sensors.external_task.ExternalTaskSensor) before the start step of the flow.
979
+ This decorator only works when a flow is scheduled on Airflow and is compiled using `airflow create`. More than one `@airflow_external_task_sensor` can be added as a flow decorators. Adding more than one decorator will ensure that `start` step starts only after all sensors finish.
980
+
981
+
982
+ Parameters
983
+ ----------
984
+ timeout : int
985
+ Time, in seconds before the task times out and fails. (Default: 3600)
986
+ poke_interval : int
987
+ Time in seconds that the job should wait in between each try. (Default: 60)
988
+ mode : str
989
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
990
+ exponential_backoff : bool
991
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
992
+ pool : str
993
+ the slot pool this task should run in,
994
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
995
+ soft_fail : bool
996
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
997
+ name : str
998
+ Name of the sensor on Airflow
999
+ description : str
1000
+ Description of sensor in the Airflow UI
1001
+ external_dag_id : str
1002
+ The dag_id that contains the task you want to wait for.
1003
+ external_task_ids : List[str]
1004
+ The list of task_ids that you want to wait for.
1005
+ If None (default value) the sensor waits for the DAG. (Default: None)
1006
+ allowed_states : List[str]
1007
+ Iterable of allowed states, (Default: ['success'])
1008
+ failed_states : List[str]
1009
+ Iterable of failed or dis-allowed states. (Default: None)
1010
+ execution_delta : datetime.timedelta
1011
+ time difference with the previous execution to look at,
1012
+ the default is the same logical date as the current task or DAG. (Default: None)
1013
+ check_existence: bool
1014
+ Set to True to check if the external task exists or check if
1015
+ the DAG to wait for exists. (Default: True)
1016
+ """
943
1017
  ...
944
1018
 
945
- @typing.overload
946
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1019
+ 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
+ """
1021
+ 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)
1022
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1023
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1024
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
1025
+ starts only after all sensors finish.
1026
+
1027
+
1028
+ Parameters
1029
+ ----------
1030
+ timeout : int
1031
+ Time, in seconds before the task times out and fails. (Default: 3600)
1032
+ poke_interval : int
1033
+ Time in seconds that the job should wait in between each try. (Default: 60)
1034
+ mode : str
1035
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1036
+ exponential_backoff : bool
1037
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1038
+ pool : str
1039
+ the slot pool this task should run in,
1040
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1041
+ soft_fail : bool
1042
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1043
+ name : str
1044
+ Name of the sensor on Airflow
1045
+ description : str
1046
+ Description of sensor in the Airflow UI
1047
+ bucket_key : Union[str, List[str]]
1048
+ The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1049
+ When it's specified as a full s3:// url, please leave `bucket_name` as None
1050
+ bucket_name : str
1051
+ Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1052
+ When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1053
+ wildcard_match : bool
1054
+ whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1055
+ aws_conn_id : str
1056
+ a reference to the s3 connection on Airflow. (Default: None)
1057
+ verify : bool
1058
+ Whether or not to verify SSL certificates for S3 connection. (Default: None)
1059
+ """
947
1060
  ...
948
1061
 
949
- 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):
1062
+ def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
950
1063
  """
951
- Specifies the PyPI packages for the step.
1064
+ Specifies what flows belong to the same project.
952
1065
 
953
- Information in this decorator will augment any
954
- attributes set in the `@pyi_base` flow-level decorator. Hence,
955
- you can use `@pypi_base` to set packages required by all
956
- steps and use `@pypi` to specify step-specific overrides.
1066
+ A project-specific namespace is created for all flows that
1067
+ use the same `@project(name)`.
957
1068
 
958
1069
 
959
1070
  Parameters
960
1071
  ----------
961
- packages : Dict[str, str], default: {}
962
- Packages to use for this step. The key is the name of the package
963
- and the value is the version to use.
964
- python : str, optional, default: None
965
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
966
- that the version used will correspond to the version of the Python interpreter used to start the run.
1072
+ name : str
1073
+ Project name. Make sure that the name is unique amongst all
1074
+ projects that use the same production scheduler. The name may
1075
+ contain only lowercase alphanumeric characters and underscores.
1076
+
1077
+ branch : Optional[str], default None
1078
+ The branch to use. If not specified, the branch is set to
1079
+ `user.<username>` unless `production` is set to `True`. This can
1080
+ also be set on the command line using `--branch` as a top-level option.
1081
+ It is an error to specify `branch` in the decorator and on the command line.
1082
+
1083
+ production : bool, default False
1084
+ Whether or not the branch is the production branch. This can also be set on the
1085
+ command line using `--production` as a top-level option. It is an error to specify
1086
+ `production` in the decorator and on the command line.
1087
+ The project branch name will be:
1088
+ - if `branch` is specified:
1089
+ - if `production` is True: `prod.<branch>`
1090
+ - if `production` is False: `test.<branch>`
1091
+ - if `branch` is not specified:
1092
+ - if `production` is True: `prod`
1093
+ - if `production` is False: `user.<username>`
967
1094
  """
968
1095
  ...
969
1096
 
970
1097
  @typing.overload
971
- def trigger(*, event: typing.Union[str, typing.Dict[str, typing.Any], None] = None, events: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], options: typing.Dict[str, typing.Any] = {}) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1098
+ def trigger_on_finish(*, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
972
1099
  """
973
- Specifies the event(s) that this flow depends on.
1100
+ Specifies the flow(s) that this flow depends on.
974
1101
 
975
1102
  ```
976
- @trigger(event='foo')
1103
+ @trigger_on_finish(flow='FooFlow')
977
1104
  ```
978
1105
  or
979
1106
  ```
980
- @trigger(events=['foo', 'bar'])
1107
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
981
1108
  ```
1109
+ This decorator respects the @project decorator and triggers the flow
1110
+ when upstream runs within the same namespace complete successfully
982
1111
 
983
- Additionally, you can specify the parameter mappings
984
- to map event payload to Metaflow parameters for the flow.
1112
+ Additionally, you can specify project aware upstream flow dependencies
1113
+ by specifying the fully qualified project_flow_name.
985
1114
  ```
986
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1115
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
987
1116
  ```
988
1117
  or
989
1118
  ```
990
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
991
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1119
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
992
1120
  ```
993
1121
 
994
- 'parameters' can also be a list of strings and tuples like so:
995
- ```
996
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
997
- ```
998
- This is equivalent to:
1122
+ You can also specify just the project or project branch (other values will be
1123
+ inferred from the current project or project branch):
999
1124
  ```
1000
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1125
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1001
1126
  ```
1002
1127
 
1128
+ Note that `branch` is typically one of:
1129
+ - `prod`
1130
+ - `user.bob`
1131
+ - `test.my_experiment`
1132
+ - `prod.staging`
1133
+
1003
1134
 
1004
1135
  Parameters
1005
1136
  ----------
1006
- event : Union[str, Dict[str, Any]], optional, default None
1007
- Event dependency for this flow.
1008
- events : List[Union[str, Dict[str, Any]]], default []
1009
- Events dependency for this flow.
1137
+ flow : Union[str, Dict[str, str]], optional, default None
1138
+ Upstream flow dependency for this flow.
1139
+ flows : List[Union[str, Dict[str, str]]], default []
1140
+ Upstream flow dependencies for this flow.
1010
1141
  options : Dict[str, Any], default {}
1011
1142
  Backend-specific configuration for tuning eventing behavior.
1012
1143
  """
1013
1144
  ...
1014
1145
 
1015
1146
  @typing.overload
1016
- def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1147
+ def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1017
1148
  ...
1018
1149
 
1019
- def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: typing.Union[str, typing.Dict[str, typing.Any], None] = None, events: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], options: typing.Dict[str, typing.Any] = {}):
1150
+ def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}):
1020
1151
  """
1021
- Specifies the event(s) that this flow depends on.
1152
+ Specifies the flow(s) that this flow depends on.
1022
1153
 
1023
1154
  ```
1024
- @trigger(event='foo')
1155
+ @trigger_on_finish(flow='FooFlow')
1025
1156
  ```
1026
1157
  or
1027
1158
  ```
1028
- @trigger(events=['foo', 'bar'])
1159
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1029
1160
  ```
1161
+ This decorator respects the @project decorator and triggers the flow
1162
+ when upstream runs within the same namespace complete successfully
1030
1163
 
1031
- Additionally, you can specify the parameter mappings
1032
- to map event payload to Metaflow parameters for the flow.
1164
+ Additionally, you can specify project aware upstream flow dependencies
1165
+ by specifying the fully qualified project_flow_name.
1033
1166
  ```
1034
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1167
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1035
1168
  ```
1036
1169
  or
1037
1170
  ```
1038
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1039
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1171
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1040
1172
  ```
1041
1173
 
1042
- 'parameters' can also be a list of strings and tuples like so:
1043
- ```
1044
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1045
- ```
1046
- This is equivalent to:
1174
+ You can also specify just the project or project branch (other values will be
1175
+ inferred from the current project or project branch):
1047
1176
  ```
1048
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1177
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1049
1178
  ```
1050
1179
 
1180
+ Note that `branch` is typically one of:
1181
+ - `prod`
1182
+ - `user.bob`
1183
+ - `test.my_experiment`
1184
+ - `prod.staging`
1185
+
1051
1186
 
1052
1187
  Parameters
1053
1188
  ----------
1054
- event : Union[str, Dict[str, Any]], optional, default None
1055
- Event dependency for this flow.
1056
- events : List[Union[str, Dict[str, Any]]], default []
1057
- Events dependency for this flow.
1189
+ flow : Union[str, Dict[str, str]], optional, default None
1190
+ Upstream flow dependency for this flow.
1191
+ flows : List[Union[str, Dict[str, str]]], default []
1192
+ Upstream flow dependencies for this flow.
1058
1193
  options : Dict[str, Any], default {}
1059
1194
  Backend-specific configuration for tuning eventing behavior.
1060
1195
  """
1061
1196
  ...
1062
1197
 
1063
1198
  @typing.overload
1064
- def trigger_on_finish(*, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1199
+ def trigger(*, event: typing.Union[str, typing.Dict[str, typing.Any], None] = None, events: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], options: typing.Dict[str, typing.Any] = {}) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1065
1200
  """
1066
- Specifies the flow(s) that this flow depends on.
1201
+ Specifies the event(s) that this flow depends on.
1067
1202
 
1068
1203
  ```
1069
- @trigger_on_finish(flow='FooFlow')
1204
+ @trigger(event='foo')
1070
1205
  ```
1071
1206
  or
1072
1207
  ```
1073
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1208
+ @trigger(events=['foo', 'bar'])
1074
1209
  ```
1075
- This decorator respects the @project decorator and triggers the flow
1076
- when upstream runs within the same namespace complete successfully
1077
1210
 
1078
- Additionally, you can specify project aware upstream flow dependencies
1079
- by specifying the fully qualified project_flow_name.
1211
+ Additionally, you can specify the parameter mappings
1212
+ to map event payload to Metaflow parameters for the flow.
1080
1213
  ```
1081
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1214
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1082
1215
  ```
1083
1216
  or
1084
1217
  ```
1085
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1218
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1219
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1086
1220
  ```
1087
1221
 
1088
- You can also specify just the project or project branch (other values will be
1089
- inferred from the current project or project branch):
1222
+ 'parameters' can also be a list of strings and tuples like so:
1223
+ ```
1224
+ @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1090
1225
  ```
1091
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1226
+ This is equivalent to:
1227
+ ```
1228
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1092
1229
  ```
1093
-
1094
- Note that `branch` is typically one of:
1095
- - `prod`
1096
- - `user.bob`
1097
- - `test.my_experiment`
1098
- - `prod.staging`
1099
1230
 
1100
1231
 
1101
1232
  Parameters
1102
1233
  ----------
1103
- flow : Union[str, Dict[str, str]], optional, default None
1104
- Upstream flow dependency for this flow.
1105
- flows : List[Union[str, Dict[str, str]]], default []
1106
- Upstream flow dependencies for this flow.
1234
+ event : Union[str, Dict[str, Any]], optional, default None
1235
+ Event dependency for this flow.
1236
+ events : List[Union[str, Dict[str, Any]]], default []
1237
+ Events dependency for this flow.
1107
1238
  options : Dict[str, Any], default {}
1108
1239
  Backend-specific configuration for tuning eventing behavior.
1109
1240
  """
1110
1241
  ...
1111
1242
 
1112
1243
  @typing.overload
1113
- def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1244
+ def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1114
1245
  ...
1115
1246
 
1116
- def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}):
1247
+ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: typing.Union[str, typing.Dict[str, typing.Any], None] = None, events: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], options: typing.Dict[str, typing.Any] = {}):
1117
1248
  """
1118
- Specifies the flow(s) that this flow depends on.
1249
+ Specifies the event(s) that this flow depends on.
1119
1250
 
1120
1251
  ```
1121
- @trigger_on_finish(flow='FooFlow')
1252
+ @trigger(event='foo')
1122
1253
  ```
1123
1254
  or
1124
1255
  ```
1125
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1256
+ @trigger(events=['foo', 'bar'])
1126
1257
  ```
1127
- This decorator respects the @project decorator and triggers the flow
1128
- when upstream runs within the same namespace complete successfully
1129
1258
 
1130
- Additionally, you can specify project aware upstream flow dependencies
1131
- by specifying the fully qualified project_flow_name.
1259
+ Additionally, you can specify the parameter mappings
1260
+ to map event payload to Metaflow parameters for the flow.
1132
1261
  ```
1133
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1262
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1134
1263
  ```
1135
1264
  or
1136
1265
  ```
1137
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1266
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1267
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1138
1268
  ```
1139
1269
 
1140
- You can also specify just the project or project branch (other values will be
1141
- inferred from the current project or project branch):
1270
+ 'parameters' can also be a list of strings and tuples like so:
1142
1271
  ```
1143
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1272
+ @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1273
+ ```
1274
+ This is equivalent to:
1275
+ ```
1276
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1144
1277
  ```
1145
-
1146
- Note that `branch` is typically one of:
1147
- - `prod`
1148
- - `user.bob`
1149
- - `test.my_experiment`
1150
- - `prod.staging`
1151
1278
 
1152
1279
 
1153
1280
  Parameters
1154
1281
  ----------
1155
- flow : Union[str, Dict[str, str]], optional, default None
1156
- Upstream flow dependency for this flow.
1157
- flows : List[Union[str, Dict[str, str]]], default []
1158
- Upstream flow dependencies for this flow.
1282
+ event : Union[str, Dict[str, Any]], optional, default None
1283
+ Event dependency for this flow.
1284
+ events : List[Union[str, Dict[str, Any]]], default []
1285
+ Events dependency for this flow.
1159
1286
  options : Dict[str, Any], default {}
1160
1287
  Backend-specific configuration for tuning eventing behavior.
1161
1288
  """
1162
1289
  ...
1163
1290
 
1164
1291
  @typing.overload
1165
- 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]]:
1292
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1166
1293
  """
1167
- Specifies the Conda environment for all steps of the flow.
1168
-
1169
- Use `@conda_base` to set common libraries required by all
1170
- steps and use `@conda` to specify step-specific additions.
1294
+ Specifies the PyPI packages for all steps of the flow.
1171
1295
 
1296
+ Use `@pypi_base` to set common packages required by all
1297
+ steps and use `@pypi` to specify step-specific overrides.
1172
1298
 
1173
1299
  Parameters
1174
1300
  ----------
1175
- packages : Dict[str, str], default {}
1301
+ packages : Dict[str, str], default: {}
1176
1302
  Packages to use for this flow. The key is the name of the package
1177
1303
  and the value is the version to use.
1178
- libraries : Dict[str, str], default {}
1179
- Supported for backward compatibility. When used with packages, packages will take precedence.
1180
- python : str, optional, default None
1304
+ python : str, optional, default: None
1181
1305
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
1182
1306
  that the version used will correspond to the version of the Python interpreter used to start the run.
1183
- disabled : bool, default False
1184
- If set to True, disables Conda.
1185
1307
  """
1186
1308
  ...
1187
1309
 
1188
1310
  @typing.overload
1189
- def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1311
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1190
1312
  ...
1191
1313
 
1192
- 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):
1314
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1193
1315
  """
1194
- Specifies the Conda environment for all steps of the flow.
1195
-
1196
- Use `@conda_base` to set common libraries required by all
1197
- steps and use `@conda` to specify step-specific additions.
1316
+ Specifies the PyPI packages for all steps of the flow.
1198
1317
 
1318
+ Use `@pypi_base` to set common packages required by all
1319
+ steps and use `@pypi` to specify step-specific overrides.
1199
1320
 
1200
1321
  Parameters
1201
1322
  ----------
1202
- packages : Dict[str, str], default {}
1323
+ packages : Dict[str, str], default: {}
1203
1324
  Packages to use for this flow. The key is the name of the package
1204
1325
  and the value is the version to use.
1205
- libraries : Dict[str, str], default {}
1206
- Supported for backward compatibility. When used with packages, packages will take precedence.
1207
- python : str, optional, default None
1326
+ python : str, optional, default: None
1208
1327
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
1209
1328
  that the version used will correspond to the version of the Python interpreter used to start the run.
1210
- disabled : bool, default False
1211
- If set to True, disables Conda.
1212
- """
1213
- ...
1214
-
1215
- 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]]:
1216
- """
1217
- 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)
1218
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1219
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1220
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
1221
- starts only after all sensors finish.
1222
-
1223
-
1224
- Parameters
1225
- ----------
1226
- timeout : int
1227
- Time, in seconds before the task times out and fails. (Default: 3600)
1228
- poke_interval : int
1229
- Time in seconds that the job should wait in between each try. (Default: 60)
1230
- mode : str
1231
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1232
- exponential_backoff : bool
1233
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1234
- pool : str
1235
- the slot pool this task should run in,
1236
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1237
- soft_fail : bool
1238
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1239
- name : str
1240
- Name of the sensor on Airflow
1241
- description : str
1242
- Description of sensor in the Airflow UI
1243
- bucket_key : Union[str, List[str]]
1244
- The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1245
- When it's specified as a full s3:// url, please leave `bucket_name` as None
1246
- bucket_name : str
1247
- Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1248
- When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1249
- wildcard_match : bool
1250
- whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1251
- aws_conn_id : str
1252
- a reference to the s3 connection on Airflow. (Default: None)
1253
- verify : bool
1254
- Whether or not to verify SSL certificates for S3 connection. (Default: None)
1255
1329
  """
1256
1330
  ...
1257
1331
 
@@ -1306,79 +1380,54 @@ def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly:
1306
1380
  """
1307
1381
  ...
1308
1382
 
1309
- def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1310
- """
1311
- Specifies what flows belong to the same project.
1312
-
1313
- A project-specific namespace is created for all flows that
1314
- use the same `@project(name)`.
1315
-
1316
-
1317
- Parameters
1318
- ----------
1319
- name : str
1320
- Project name. Make sure that the name is unique amongst all
1321
- projects that use the same production scheduler. The name may
1322
- contain only lowercase alphanumeric characters and underscores.
1323
-
1324
- branch : Optional[str], default None
1325
- The branch to use. If not specified, the branch is set to
1326
- `user.<username>` unless `production` is set to `True`. This can
1327
- also be set on the command line using `--branch` as a top-level option.
1328
- It is an error to specify `branch` in the decorator and on the command line.
1329
-
1330
- production : bool, default False
1331
- Whether or not the branch is the production branch. This can also be set on the
1332
- command line using `--production` as a top-level option. It is an error to specify
1333
- `production` in the decorator and on the command line.
1334
- The project branch name will be:
1335
- - if `branch` is specified:
1336
- - if `production` is True: `prod.<branch>`
1337
- - if `production` is False: `test.<branch>`
1338
- - if `branch` is not specified:
1339
- - if `production` is True: `prod`
1340
- - if `production` is False: `user.<username>`
1341
- """
1342
- ...
1343
-
1344
1383
  @typing.overload
1345
- def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1384
+ 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]]:
1346
1385
  """
1347
- Specifies the PyPI packages for all steps of the flow.
1386
+ Specifies the Conda environment for all steps of the flow.
1387
+
1388
+ Use `@conda_base` to set common libraries required by all
1389
+ steps and use `@conda` to specify step-specific additions.
1348
1390
 
1349
- Use `@pypi_base` to set common packages required by all
1350
- steps and use `@pypi` to specify step-specific overrides.
1351
1391
 
1352
1392
  Parameters
1353
1393
  ----------
1354
- packages : Dict[str, str], default: {}
1394
+ packages : Dict[str, str], default {}
1355
1395
  Packages to use for this flow. The key is the name of the package
1356
1396
  and the value is the version to use.
1357
- python : str, optional, default: None
1397
+ libraries : Dict[str, str], default {}
1398
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1399
+ python : str, optional, default None
1358
1400
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
1359
1401
  that the version used will correspond to the version of the Python interpreter used to start the run.
1402
+ disabled : bool, default False
1403
+ If set to True, disables Conda.
1360
1404
  """
1361
1405
  ...
1362
1406
 
1363
1407
  @typing.overload
1364
- def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1408
+ def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1365
1409
  ...
1366
1410
 
1367
- def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1411
+ 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):
1368
1412
  """
1369
- Specifies the PyPI packages for all steps of the flow.
1413
+ Specifies the Conda environment for all steps of the flow.
1414
+
1415
+ Use `@conda_base` to set common libraries required by all
1416
+ steps and use `@conda` to specify step-specific additions.
1370
1417
 
1371
- Use `@pypi_base` to set common packages required by all
1372
- steps and use `@pypi` to specify step-specific overrides.
1373
1418
 
1374
1419
  Parameters
1375
1420
  ----------
1376
- packages : Dict[str, str], default: {}
1421
+ packages : Dict[str, str], default {}
1377
1422
  Packages to use for this flow. The key is the name of the package
1378
1423
  and the value is the version to use.
1379
- python : str, optional, default: None
1424
+ libraries : Dict[str, str], default {}
1425
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1426
+ python : str, optional, default None
1380
1427
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
1381
1428
  that the version used will correspond to the version of the Python interpreter used to start the run.
1429
+ disabled : bool, default False
1430
+ If set to True, disables Conda.
1382
1431
  """
1383
1432
  ...
1384
1433
 
@@ -1496,48 +1545,5 @@ def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None)
1496
1545
  """
1497
1546
  ...
1498
1547
 
1499
- def airflow_external_task_sensor(*, timeout: int, poke_interval: int, mode: str, exponential_backoff: bool, pool: str, soft_fail: bool, name: str, description: str, external_dag_id: str, external_task_ids: typing.List[str], allowed_states: typing.List[str], failed_states: typing.List[str], execution_delta: "datetime.timedelta", check_existence: bool) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1500
- """
1501
- The `@airflow_external_task_sensor` decorator attaches a Airflow [ExternalTaskSensor](https://airflow.apache.org/docs/apache-airflow/stable/_api/airflow/sensors/external_task/index.html#airflow.sensors.external_task.ExternalTaskSensor) before the start step of the flow.
1502
- This decorator only works when a flow is scheduled on Airflow and is compiled using `airflow create`. More than one `@airflow_external_task_sensor` can be added as a flow decorators. Adding more than one decorator will ensure that `start` step starts only after all sensors finish.
1503
-
1504
-
1505
- Parameters
1506
- ----------
1507
- timeout : int
1508
- Time, in seconds before the task times out and fails. (Default: 3600)
1509
- poke_interval : int
1510
- Time in seconds that the job should wait in between each try. (Default: 60)
1511
- mode : str
1512
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1513
- exponential_backoff : bool
1514
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1515
- pool : str
1516
- the slot pool this task should run in,
1517
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1518
- soft_fail : bool
1519
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1520
- name : str
1521
- Name of the sensor on Airflow
1522
- description : str
1523
- Description of sensor in the Airflow UI
1524
- external_dag_id : str
1525
- The dag_id that contains the task you want to wait for.
1526
- external_task_ids : List[str]
1527
- The list of task_ids that you want to wait for.
1528
- If None (default value) the sensor waits for the DAG. (Default: None)
1529
- allowed_states : List[str]
1530
- Iterable of allowed states, (Default: ['success'])
1531
- failed_states : List[str]
1532
- Iterable of failed or dis-allowed states. (Default: None)
1533
- execution_delta : datetime.timedelta
1534
- time difference with the previous execution to look at,
1535
- the default is the same logical date as the current task or DAG. (Default: None)
1536
- check_existence: bool
1537
- Set to True to check if the external task exists or check if
1538
- the DAG to wait for exists. (Default: True)
1539
- """
1540
- ...
1541
-
1542
1548
  pkg_name: str
1543
1549