ob-metaflow-stubs 3.9__py2.py3-none-any.whl → 4.1__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 (142) hide show
  1. metaflow-stubs/__init__.pyi +577 -577
  2. metaflow-stubs/cards.pyi +11 -5
  3. metaflow-stubs/cli.pyi +2 -2
  4. metaflow-stubs/client/__init__.pyi +3 -3
  5. metaflow-stubs/client/core.pyi +5 -5
  6. metaflow-stubs/client/filecache.pyi +2 -2
  7. metaflow-stubs/clone_util.pyi +2 -2
  8. metaflow-stubs/events.pyi +2 -2
  9. metaflow-stubs/exception.pyi +2 -2
  10. metaflow-stubs/flowspec.pyi +5 -5
  11. metaflow-stubs/generated_for.txt +1 -1
  12. metaflow-stubs/includefile.pyi +4 -4
  13. metaflow-stubs/metadata/metadata.pyi +3 -3
  14. metaflow-stubs/metadata/util.pyi +2 -2
  15. metaflow-stubs/metaflow_config.pyi +4 -2
  16. metaflow-stubs/metaflow_current.pyi +17 -17
  17. metaflow-stubs/mflog/mflog.pyi +2 -2
  18. metaflow-stubs/multicore_utils.pyi +2 -2
  19. metaflow-stubs/parameters.pyi +4 -4
  20. metaflow-stubs/plugins/__init__.pyi +8 -2
  21. metaflow-stubs/plugins/airflow/__init__.pyi +2 -2
  22. metaflow-stubs/plugins/airflow/airflow.pyi +3 -3
  23. metaflow-stubs/plugins/airflow/airflow_cli.pyi +3 -3
  24. metaflow-stubs/plugins/airflow/airflow_decorator.pyi +2 -2
  25. metaflow-stubs/plugins/airflow/airflow_utils.pyi +2 -2
  26. metaflow-stubs/plugins/airflow/exception.pyi +2 -2
  27. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +2 -2
  28. metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi +3 -3
  29. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +3 -3
  30. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +3 -3
  31. metaflow-stubs/plugins/argo/__init__.pyi +2 -2
  32. metaflow-stubs/plugins/argo/argo_client.pyi +2 -2
  33. metaflow-stubs/plugins/argo/argo_events.pyi +2 -2
  34. metaflow-stubs/plugins/argo/argo_workflows.pyi +7 -3
  35. metaflow-stubs/plugins/argo/argo_workflows_cli.pyi +4 -4
  36. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +4 -4
  37. metaflow-stubs/plugins/aws/__init__.pyi +2 -2
  38. metaflow-stubs/plugins/aws/aws_client.pyi +2 -2
  39. metaflow-stubs/plugins/aws/aws_utils.pyi +2 -2
  40. metaflow-stubs/plugins/aws/batch/__init__.pyi +2 -2
  41. metaflow-stubs/plugins/aws/batch/batch.pyi +2 -2
  42. metaflow-stubs/plugins/aws/batch/batch_cli.pyi +2 -2
  43. metaflow-stubs/plugins/aws/batch/batch_client.pyi +2 -2
  44. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +2 -2
  45. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +2 -2
  46. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +3 -3
  47. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +2 -2
  48. metaflow-stubs/plugins/aws/step_functions/dynamo_db_client.pyi +2 -2
  49. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +2 -2
  50. metaflow-stubs/plugins/aws/step_functions/production_token.pyi +2 -2
  51. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +2 -2
  52. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +2 -2
  53. metaflow-stubs/plugins/aws/step_functions/step_functions_cli.pyi +4 -4
  54. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +2 -2
  55. metaflow-stubs/plugins/aws/step_functions/step_functions_decorator.pyi +2 -2
  56. metaflow-stubs/plugins/azure/__init__.pyi +2 -2
  57. metaflow-stubs/plugins/azure/azure_credential.pyi +2 -2
  58. metaflow-stubs/plugins/azure/azure_exceptions.pyi +2 -2
  59. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +3 -3
  60. metaflow-stubs/plugins/azure/azure_utils.pyi +2 -2
  61. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +2 -2
  62. metaflow-stubs/plugins/azure/includefile_support.pyi +2 -2
  63. metaflow-stubs/plugins/cards/__init__.pyi +2 -2
  64. metaflow-stubs/plugins/cards/card_cli.pyi +4 -4
  65. metaflow-stubs/plugins/cards/card_client.pyi +3 -3
  66. metaflow-stubs/plugins/cards/card_creator.pyi +2 -2
  67. metaflow-stubs/plugins/cards/card_datastore.pyi +2 -2
  68. metaflow-stubs/plugins/cards/card_decorator.pyi +2 -2
  69. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +2 -2
  70. metaflow-stubs/plugins/cards/card_modules/basic.pyi +9 -3
  71. metaflow-stubs/plugins/cards/card_modules/card.pyi +2 -2
  72. metaflow-stubs/plugins/cards/card_modules/chevron/__init__.pyi +2 -2
  73. metaflow-stubs/plugins/cards/card_modules/chevron/main.pyi +2 -2
  74. metaflow-stubs/plugins/cards/card_modules/chevron/metadata.pyi +2 -2
  75. metaflow-stubs/plugins/cards/card_modules/chevron/renderer.pyi +2 -2
  76. metaflow-stubs/plugins/cards/card_modules/chevron/tokenizer.pyi +2 -2
  77. metaflow-stubs/plugins/cards/card_modules/components.pyi +4 -4
  78. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +2 -2
  79. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +2 -2
  80. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +3 -3
  81. metaflow-stubs/plugins/cards/card_resolver.pyi +2 -2
  82. metaflow-stubs/plugins/cards/component_serializer.pyi +3 -3
  83. metaflow-stubs/plugins/cards/exception.pyi +2 -2
  84. metaflow-stubs/plugins/catch_decorator.pyi +3 -3
  85. metaflow-stubs/plugins/datatools/__init__.pyi +3 -3
  86. metaflow-stubs/plugins/datatools/local.pyi +2 -2
  87. metaflow-stubs/plugins/datatools/s3/__init__.pyi +3 -3
  88. metaflow-stubs/plugins/datatools/s3/s3.pyi +6 -6
  89. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +2 -2
  90. metaflow-stubs/plugins/datatools/s3/s3util.pyi +2 -2
  91. metaflow-stubs/plugins/debug_logger.pyi +2 -2
  92. metaflow-stubs/plugins/debug_monitor.pyi +2 -2
  93. metaflow-stubs/plugins/environment_decorator.pyi +2 -2
  94. metaflow-stubs/plugins/events_decorator.pyi +2 -2
  95. metaflow-stubs/plugins/frameworks/__init__.pyi +2 -2
  96. metaflow-stubs/plugins/frameworks/pytorch.pyi +3 -3
  97. metaflow-stubs/plugins/gcp/__init__.pyi +2 -2
  98. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +3 -3
  99. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +2 -2
  100. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +2 -2
  101. metaflow-stubs/plugins/gcp/gs_utils.pyi +2 -2
  102. metaflow-stubs/plugins/gcp/includefile_support.pyi +2 -2
  103. metaflow-stubs/plugins/kubernetes/__init__.pyi +2 -2
  104. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +2 -2
  105. metaflow-stubs/plugins/kubernetes/kubernetes_cli.pyi +3 -3
  106. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +2 -2
  107. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +2 -2
  108. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +2 -2
  109. metaflow-stubs/plugins/logs_cli.pyi +2 -2
  110. metaflow-stubs/plugins/package_cli.pyi +2 -2
  111. metaflow-stubs/plugins/parallel_decorator.pyi +2 -2
  112. metaflow-stubs/plugins/perimeters.pyi +2 -2
  113. metaflow-stubs/plugins/project_decorator.pyi +2 -2
  114. metaflow-stubs/plugins/pypi/__init__.pyi +2 -2
  115. metaflow-stubs/plugins/pypi/conda_decorator.pyi +2 -2
  116. metaflow-stubs/plugins/pypi/conda_environment.pyi +5 -5
  117. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +2 -2
  118. metaflow-stubs/plugins/pypi/pypi_environment.pyi +3 -3
  119. metaflow-stubs/plugins/pypi/utils.pyi +2 -2
  120. metaflow-stubs/plugins/resources_decorator.pyi +2 -2
  121. metaflow-stubs/plugins/retry_decorator.pyi +2 -2
  122. metaflow-stubs/plugins/secrets/__init__.pyi +2 -2
  123. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +2 -2
  124. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +2 -2
  125. metaflow-stubs/plugins/storage_executor.pyi +2 -2
  126. metaflow-stubs/plugins/tag_cli.pyi +3 -3
  127. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +2 -2
  128. metaflow-stubs/plugins/timeout_decorator.pyi +3 -3
  129. metaflow-stubs/procpoll.pyi +2 -2
  130. metaflow-stubs/profilers/__init__.pyi +2 -2
  131. metaflow-stubs/pylint_wrapper.pyi +2 -2
  132. metaflow-stubs/runner/__init__.pyi +2 -2
  133. metaflow-stubs/runner/metaflow_runner.pyi +4 -4
  134. metaflow-stubs/runner/nbrun.pyi +2 -2
  135. metaflow-stubs/runner/subprocess_manager.pyi +2 -2
  136. metaflow-stubs/tagging_util.pyi +2 -2
  137. metaflow-stubs/tuple_util.pyi +2 -2
  138. {ob_metaflow_stubs-3.9.dist-info → ob_metaflow_stubs-4.1.dist-info}/METADATA +1 -1
  139. ob_metaflow_stubs-4.1.dist-info/RECORD +142 -0
  140. ob_metaflow_stubs-3.9.dist-info/RECORD +0 -142
  141. {ob_metaflow_stubs-3.9.dist-info → ob_metaflow_stubs-4.1.dist-info}/WHEEL +0 -0
  142. {ob_metaflow_stubs-3.9.dist-info → ob_metaflow_stubs-4.1.dist-info}/top_level.txt +0 -0
@@ -1,25 +1,25 @@
1
1
  ##################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
- # MF version: 2.12.4.2+ob(v1) #
4
- # Generated on 2024-06-19T18:39:56.270719 #
3
+ # MF version: 2.12.6.1+ob(v1) #
4
+ # Generated on 2024-07-03T18:20:45.077392 #
5
5
  ##################################################################################
6
6
 
7
7
  from __future__ import annotations
8
8
 
9
9
  import typing
10
10
  if typing.TYPE_CHECKING:
11
- import metaflow.plugins.datatools.s3.s3
12
- import metaflow._vendor.click.types
13
- import typing
11
+ import metaflow.flowspec
12
+ import metaflow.datastore.inputs
14
13
  import metaflow.events
14
+ import metaflow._vendor.click.types
15
15
  import metaflow.metaflow_current
16
- import metaflow.datastore.inputs
17
- import metaflow.client.core
18
- import metaflow.runner.metaflow_runner
19
16
  import metaflow.parameters
20
- import metaflow.flowspec
21
- import datetime
17
+ import typing
18
+ import metaflow.runner.metaflow_runner
22
19
  import io
20
+ import metaflow.client.core
21
+ import metaflow.plugins.datatools.s3.s3
22
+ import datetime
23
23
  FlowSpecDerived = typing.TypeVar("FlowSpecDerived", bound="FlowSpec", contravariant=False, covariant=False)
24
24
  StepFlag = typing.NewType("StepFlag", bool)
25
25
 
@@ -727,63 +727,6 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
727
727
  """
728
728
  ...
729
729
 
730
- @typing.overload
731
- 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]]]:
732
- """
733
- Specifies a timeout for your step.
734
-
735
- This decorator is useful if this step may hang indefinitely.
736
-
737
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
738
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
739
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
740
-
741
- Note that all the values specified in parameters are added together so if you specify
742
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
743
-
744
- Parameters
745
- ----------
746
- seconds : int, default 0
747
- Number of seconds to wait prior to timing out.
748
- minutes : int, default 0
749
- Number of minutes to wait prior to timing out.
750
- hours : int, default 0
751
- Number of hours to wait prior to timing out.
752
- """
753
- ...
754
-
755
- @typing.overload
756
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
757
- ...
758
-
759
- @typing.overload
760
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
761
- ...
762
-
763
- 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):
764
- """
765
- Specifies a timeout for your step.
766
-
767
- This decorator is useful if this step may hang indefinitely.
768
-
769
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
770
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
771
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
772
-
773
- Note that all the values specified in parameters are added together so if you specify
774
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
775
-
776
- Parameters
777
- ----------
778
- seconds : int, default 0
779
- Number of seconds to wait prior to timing out.
780
- minutes : int, default 0
781
- Number of minutes to wait prior to timing out.
782
- hours : int, default 0
783
- Number of hours to wait prior to timing out.
784
- """
785
- ...
786
-
787
730
  @typing.overload
788
731
  def catch(*, var: typing.Optional[str] = None, print_exception: bool = True) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
789
732
  """
@@ -833,6 +776,65 @@ def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
833
776
  """
834
777
  ...
835
778
 
779
+ 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, namespace: str = "METAFLOW_KUBERNETES_NAMESPACE", gpu: typing.Optional[int] = None, gpu_vendor: str = "KUBERNETES_GPU_VENDOR", tolerations: typing.List[str] = [], 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) -> 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]]]:
780
+ """
781
+ Specifies that this step should execute on Kubernetes.
782
+
783
+ Parameters
784
+ ----------
785
+ cpu : int, default 1
786
+ Number of CPUs required for this step. If `@resources` is
787
+ also present, the maximum value from all decorators is used.
788
+ memory : int, default 4096
789
+ Memory size (in MB) required for this step. If
790
+ `@resources` is also present, the maximum value from all decorators is
791
+ used.
792
+ disk : int, default 10240
793
+ Disk size (in MB) required for this step. If
794
+ `@resources` is also present, the maximum value from all decorators is
795
+ used.
796
+ image : str, optional, default None
797
+ Docker image to use when launching on Kubernetes. If not specified, and
798
+ METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
799
+ not, a default Docker image mapping to the current version of Python is used.
800
+ image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
801
+ If given, the imagePullPolicy to be applied to the Docker image of the step.
802
+ service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
803
+ Kubernetes service account to use when launching pod in Kubernetes.
804
+ secrets : List[str], optional, default None
805
+ Kubernetes secrets to use when launching pod in Kubernetes. These
806
+ secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
807
+ in Metaflow configuration.
808
+ namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
809
+ Kubernetes namespace to use when launching pod in Kubernetes.
810
+ gpu : int, optional, default None
811
+ Number of GPUs required for this step. A value of zero implies that
812
+ the scheduled node should not have GPUs.
813
+ gpu_vendor : str, default KUBERNETES_GPU_VENDOR
814
+ The vendor of the GPUs to be used for this step.
815
+ tolerations : List[str], default []
816
+ The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
817
+ Kubernetes tolerations to use when launching pod in Kubernetes.
818
+ use_tmpfs : bool, default False
819
+ This enables an explicit tmpfs mount for this step.
820
+ tmpfs_tempdir : bool, default True
821
+ sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
822
+ tmpfs_size : int, optional, default: None
823
+ The value for the size (in MiB) of the tmpfs mount for this step.
824
+ This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
825
+ memory allocated for this step.
826
+ tmpfs_path : str, optional, default /metaflow_temp
827
+ Path to tmpfs mount for this step.
828
+ persistent_volume_claims : Dict[str, str], optional, default None
829
+ A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
830
+ volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
831
+ shared_memory: int, optional
832
+ Shared memory size (in MiB) required for this step
833
+ port: int, optional
834
+ Port number to specify in the Kubernetes job object
835
+ """
836
+ ...
837
+
836
838
  @typing.overload
837
839
  def card(*, type: str = "default", id: typing.Optional[str] = None, options: typing.Dict[str, typing.Any] = {}, timeout: int = 45) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
838
840
  """
@@ -885,368 +887,66 @@ def card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typ
885
887
  ...
886
888
 
887
889
  @typing.overload
888
- 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]]]:
889
- """
890
- Specifies the Conda environment for the step.
891
-
892
- Information in this decorator will augment any
893
- attributes set in the `@conda_base` flow-level decorator. Hence,
894
- you can use `@conda_base` to set packages required by all
895
- steps and use `@conda` to specify step-specific overrides.
896
-
897
- Parameters
898
- ----------
899
- packages : Dict[str, str], default {}
900
- Packages to use for this step. The key is the name of the package
901
- and the value is the version to use.
902
- libraries : Dict[str, str], default {}
903
- Supported for backward compatibility. When used with packages, packages will take precedence.
904
- python : str, optional, default None
905
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
906
- that the version used will correspond to the version of the Python interpreter used to start the run.
907
- disabled : bool, default False
908
- If set to True, disables @conda.
909
- """
910
- ...
911
-
912
- @typing.overload
913
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
914
- ...
915
-
916
- @typing.overload
917
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
918
- ...
919
-
920
- 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):
890
+ 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]]]:
921
891
  """
922
- Specifies the Conda environment for the step.
923
-
924
- Information in this decorator will augment any
925
- attributes set in the `@conda_base` flow-level decorator. Hence,
926
- you can use `@conda_base` to set packages required by all
927
- steps and use `@conda` to specify step-specific overrides.
892
+ Specifies a timeout for your step.
928
893
 
929
- Parameters
930
- ----------
931
- packages : Dict[str, str], default {}
932
- Packages to use for this step. The key is the name of the package
933
- and the value is the version to use.
934
- libraries : Dict[str, str], default {}
935
- Supported for backward compatibility. When used with packages, packages will take precedence.
936
- python : str, optional, default None
937
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
938
- that the version used will correspond to the version of the Python interpreter used to start the run.
939
- disabled : bool, default False
940
- If set to True, disables @conda.
941
- """
942
- ...
943
-
944
- @typing.overload
945
- 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]]]:
946
- """
947
- Specifies the number of times the task corresponding
948
- to a step needs to be retried.
894
+ This decorator is useful if this step may hang indefinitely.
949
895
 
950
- This decorator is useful for handling transient errors, such as networking issues.
951
- If your task contains operations that can't be retried safely, e.g. database updates,
952
- it is advisable to annotate it with `@retry(times=0)`.
896
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
897
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
898
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
953
899
 
954
- This can be used in conjunction with the `@catch` decorator. The `@catch`
955
- decorator will execute a no-op task after all retries have been exhausted,
956
- ensuring that the flow execution can continue.
900
+ Note that all the values specified in parameters are added together so if you specify
901
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
957
902
 
958
903
  Parameters
959
904
  ----------
960
- times : int, default 3
961
- Number of times to retry this task.
962
- minutes_between_retries : int, default 2
963
- Number of minutes between retries.
905
+ seconds : int, default 0
906
+ Number of seconds to wait prior to timing out.
907
+ minutes : int, default 0
908
+ Number of minutes to wait prior to timing out.
909
+ hours : int, default 0
910
+ Number of hours to wait prior to timing out.
964
911
  """
965
912
  ...
966
913
 
967
914
  @typing.overload
968
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
915
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
969
916
  ...
970
917
 
971
918
  @typing.overload
972
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
919
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
973
920
  ...
974
921
 
975
- 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):
922
+ 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):
976
923
  """
977
- Specifies the number of times the task corresponding
978
- to a step needs to be retried.
924
+ Specifies a timeout for your step.
979
925
 
980
- This decorator is useful for handling transient errors, such as networking issues.
981
- If your task contains operations that can't be retried safely, e.g. database updates,
982
- it is advisable to annotate it with `@retry(times=0)`.
926
+ This decorator is useful if this step may hang indefinitely.
983
927
 
984
- This can be used in conjunction with the `@catch` decorator. The `@catch`
985
- decorator will execute a no-op task after all retries have been exhausted,
986
- ensuring that the flow execution can continue.
928
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
929
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
930
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
931
+
932
+ Note that all the values specified in parameters are added together so if you specify
933
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
987
934
 
988
935
  Parameters
989
936
  ----------
990
- times : int, default 3
991
- Number of times to retry this task.
992
- minutes_between_retries : int, default 2
993
- Number of minutes between retries.
937
+ seconds : int, default 0
938
+ Number of seconds to wait prior to timing out.
939
+ minutes : int, default 0
940
+ Number of minutes to wait prior to timing out.
941
+ hours : int, default 0
942
+ Number of hours to wait prior to timing out.
994
943
  """
995
944
  ...
996
945
 
997
946
  @typing.overload
998
- def secrets(*, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = []) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
947
+ def batch(*, cpu: int = 1, gpu: int = 0, memory: int = 4096, image: typing.Optional[str] = None, queue: str = "METAFLOW_BATCH_JOB_QUEUE", iam_role: str = "METAFLOW_ECS_S3_ACCESS_IAM_ROLE", execution_role: str = "METAFLOW_ECS_FARGATE_EXECUTION_ROLE", shared_memory: typing.Optional[int] = None, max_swap: typing.Optional[int] = None, swappiness: typing.Optional[int] = None, use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = None, inferentia: int = 0, trainium: int = None, efa: int = 0, ephemeral_storage: int = None, log_driver: typing.Optional[str] = None, log_options: typing.Optional[typing.List[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]]]:
999
948
  """
1000
- Specifies secrets to be retrieved and injected as environment variables prior to
1001
- the execution of a step.
1002
-
1003
- Parameters
1004
- ----------
1005
- sources : List[Union[str, Dict[str, Any]]], default: []
1006
- List of secret specs, defining how the secrets are to be retrieved
1007
- """
1008
- ...
1009
-
1010
- @typing.overload
1011
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1012
- ...
1013
-
1014
- @typing.overload
1015
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1016
- ...
1017
-
1018
- def secrets(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = []):
1019
- """
1020
- Specifies secrets to be retrieved and injected as environment variables prior to
1021
- the execution of a step.
1022
-
1023
- Parameters
1024
- ----------
1025
- sources : List[Union[str, Dict[str, Any]]], default: []
1026
- List of secret specs, defining how the secrets are to be retrieved
1027
- """
1028
- ...
1029
-
1030
- @typing.overload
1031
- 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]]]:
1032
- """
1033
- Specifies the PyPI packages for the step.
1034
-
1035
- Information in this decorator will augment any
1036
- attributes set in the `@pyi_base` flow-level decorator. Hence,
1037
- you can use `@pypi_base` to set packages required by all
1038
- steps and use `@pypi` to specify step-specific overrides.
1039
-
1040
- Parameters
1041
- ----------
1042
- packages : Dict[str, str], default: {}
1043
- Packages to use for this step. The key is the name of the package
1044
- and the value is the version to use.
1045
- python : str, optional, default: None
1046
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1047
- that the version used will correspond to the version of the Python interpreter used to start the run.
1048
- """
1049
- ...
1050
-
1051
- @typing.overload
1052
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1053
- ...
1054
-
1055
- @typing.overload
1056
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1057
- ...
1058
-
1059
- 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):
1060
- """
1061
- Specifies the PyPI packages for the step.
1062
-
1063
- Information in this decorator will augment any
1064
- attributes set in the `@pyi_base` flow-level decorator. Hence,
1065
- you can use `@pypi_base` to set packages required by all
1066
- steps and use `@pypi` to specify step-specific overrides.
1067
-
1068
- Parameters
1069
- ----------
1070
- packages : Dict[str, str], default: {}
1071
- Packages to use for this step. The key is the name of the package
1072
- and the value is the version to use.
1073
- python : str, optional, default: None
1074
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1075
- that the version used will correspond to the version of the Python interpreter used to start the run.
1076
- """
1077
- ...
1078
-
1079
- 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, namespace: str = "METAFLOW_KUBERNETES_NAMESPACE", gpu: typing.Optional[int] = None, gpu_vendor: str = "KUBERNETES_GPU_VENDOR", tolerations: typing.List[str] = [], 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) -> 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]]]:
1080
- """
1081
- Specifies that this step should execute on Kubernetes.
1082
-
1083
- Parameters
1084
- ----------
1085
- cpu : int, default 1
1086
- Number of CPUs required for this step. If `@resources` is
1087
- also present, the maximum value from all decorators is used.
1088
- memory : int, default 4096
1089
- Memory size (in MB) required for this step. If
1090
- `@resources` is also present, the maximum value from all decorators is
1091
- used.
1092
- disk : int, default 10240
1093
- Disk size (in MB) required for this step. If
1094
- `@resources` is also present, the maximum value from all decorators is
1095
- used.
1096
- image : str, optional, default None
1097
- Docker image to use when launching on Kubernetes. If not specified, and
1098
- METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
1099
- not, a default Docker image mapping to the current version of Python is used.
1100
- image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
1101
- If given, the imagePullPolicy to be applied to the Docker image of the step.
1102
- service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
1103
- Kubernetes service account to use when launching pod in Kubernetes.
1104
- secrets : List[str], optional, default None
1105
- Kubernetes secrets to use when launching pod in Kubernetes. These
1106
- secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
1107
- in Metaflow configuration.
1108
- namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
1109
- Kubernetes namespace to use when launching pod in Kubernetes.
1110
- gpu : int, optional, default None
1111
- Number of GPUs required for this step. A value of zero implies that
1112
- the scheduled node should not have GPUs.
1113
- gpu_vendor : str, default KUBERNETES_GPU_VENDOR
1114
- The vendor of the GPUs to be used for this step.
1115
- tolerations : List[str], default []
1116
- The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
1117
- Kubernetes tolerations to use when launching pod in Kubernetes.
1118
- use_tmpfs : bool, default False
1119
- This enables an explicit tmpfs mount for this step.
1120
- tmpfs_tempdir : bool, default True
1121
- sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
1122
- tmpfs_size : int, optional, default: None
1123
- The value for the size (in MiB) of the tmpfs mount for this step.
1124
- This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
1125
- memory allocated for this step.
1126
- tmpfs_path : str, optional, default /metaflow_temp
1127
- Path to tmpfs mount for this step.
1128
- persistent_volume_claims : Dict[str, str], optional, default None
1129
- A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
1130
- volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
1131
- shared_memory: int, optional
1132
- Shared memory size (in MiB) required for this step
1133
- port: int, optional
1134
- Port number to specify in the Kubernetes job object
1135
- """
1136
- ...
1137
-
1138
- @typing.overload
1139
- 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]]]:
1140
- """
1141
- Specifies environment variables to be set prior to the execution of a step.
1142
-
1143
- Parameters
1144
- ----------
1145
- vars : Dict[str, str], default {}
1146
- Dictionary of environment variables to set.
1147
- """
1148
- ...
1149
-
1150
- @typing.overload
1151
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1152
- ...
1153
-
1154
- @typing.overload
1155
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1156
- ...
1157
-
1158
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
1159
- """
1160
- Specifies environment variables to be set prior to the execution of a step.
1161
-
1162
- Parameters
1163
- ----------
1164
- vars : Dict[str, str], default {}
1165
- Dictionary of environment variables to set.
1166
- """
1167
- ...
1168
-
1169
- @typing.overload
1170
- def resources(*, cpu: int = 1, gpu: int = 0, 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]]]:
1171
- """
1172
- Specifies the resources needed when executing this step.
1173
-
1174
- Use `@resources` to specify the resource requirements
1175
- independently of the specific compute layer (`@batch`, `@kubernetes`).
1176
-
1177
- You can choose the compute layer on the command line by executing e.g.
1178
- ```
1179
- python myflow.py run --with batch
1180
- ```
1181
- or
1182
- ```
1183
- python myflow.py run --with kubernetes
1184
- ```
1185
- which executes the flow on the desired system using the
1186
- requirements specified in `@resources`.
1187
-
1188
- Parameters
1189
- ----------
1190
- cpu : int, default 1
1191
- Number of CPUs required for this step.
1192
- gpu : int, default 0
1193
- Number of GPUs required for this step.
1194
- disk : int, optional, default None
1195
- Disk size (in MB) required for this step. Only applies on Kubernetes.
1196
- memory : int, default 4096
1197
- Memory size (in MB) required for this step.
1198
- shared_memory : int, optional, default None
1199
- The value for the size (in MiB) of the /dev/shm volume for this step.
1200
- This parameter maps to the `--shm-size` option in Docker.
1201
- """
1202
- ...
1203
-
1204
- @typing.overload
1205
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1206
- ...
1207
-
1208
- @typing.overload
1209
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1210
- ...
1211
-
1212
- def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, cpu: int = 1, gpu: int = 0, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None):
1213
- """
1214
- Specifies the resources needed when executing this step.
1215
-
1216
- Use `@resources` to specify the resource requirements
1217
- independently of the specific compute layer (`@batch`, `@kubernetes`).
1218
-
1219
- You can choose the compute layer on the command line by executing e.g.
1220
- ```
1221
- python myflow.py run --with batch
1222
- ```
1223
- or
1224
- ```
1225
- python myflow.py run --with kubernetes
1226
- ```
1227
- which executes the flow on the desired system using the
1228
- requirements specified in `@resources`.
1229
-
1230
- Parameters
1231
- ----------
1232
- cpu : int, default 1
1233
- Number of CPUs required for this step.
1234
- gpu : int, default 0
1235
- Number of GPUs required for this step.
1236
- disk : int, optional, default None
1237
- Disk size (in MB) required for this step. Only applies on Kubernetes.
1238
- memory : int, default 4096
1239
- Memory size (in MB) required for this step.
1240
- shared_memory : int, optional, default None
1241
- The value for the size (in MiB) of the /dev/shm volume for this step.
1242
- This parameter maps to the `--shm-size` option in Docker.
1243
- """
1244
- ...
1245
-
1246
- @typing.overload
1247
- def batch(*, cpu: int = 1, gpu: int = 0, memory: int = 4096, image: typing.Optional[str] = None, queue: str = "METAFLOW_BATCH_JOB_QUEUE", iam_role: str = "METAFLOW_ECS_S3_ACCESS_IAM_ROLE", execution_role: str = "METAFLOW_ECS_FARGATE_EXECUTION_ROLE", shared_memory: typing.Optional[int] = None, max_swap: typing.Optional[int] = None, swappiness: typing.Optional[int] = None, use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = None, inferentia: int = 0, trainium: int = None, efa: int = 0, ephemeral_storage: int = None, log_driver: typing.Optional[str] = None, log_options: typing.Optional[typing.List[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]]]:
1248
- """
1249
- Specifies that this step should execute on [AWS Batch](https://aws.amazon.com/batch/).
949
+ Specifies that this step should execute on [AWS Batch](https://aws.amazon.com/batch/).
1250
950
 
1251
951
  Parameters
1252
952
  ----------
@@ -1390,120 +1090,183 @@ def batch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
1390
1090
  """
1391
1091
  ...
1392
1092
 
1393
- 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]]:
1093
+ @typing.overload
1094
+ 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]]]:
1394
1095
  """
1395
- 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)
1396
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1397
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1398
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
1399
- starts only after all sensors finish.
1096
+ Specifies the number of times the task corresponding
1097
+ to a step needs to be retried.
1400
1098
 
1401
- Parameters
1402
- ----------
1403
- timeout : int
1404
- Time, in seconds before the task times out and fails. (Default: 3600)
1405
- poke_interval : int
1406
- Time in seconds that the job should wait in between each try. (Default: 60)
1407
- mode : str
1408
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1409
- exponential_backoff : bool
1410
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1411
- pool : str
1412
- the slot pool this task should run in,
1413
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1414
- soft_fail : bool
1415
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1416
- name : str
1417
- Name of the sensor on Airflow
1418
- description : str
1419
- Description of sensor in the Airflow UI
1420
- bucket_key : Union[str, List[str]]
1421
- The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1422
- When it's specified as a full s3:// url, please leave `bucket_name` as None
1423
- bucket_name : str
1424
- Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1425
- When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1426
- wildcard_match : bool
1427
- whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1428
- aws_conn_id : str
1429
- a reference to the s3 connection on Airflow. (Default: None)
1430
- verify : bool
1431
- Whether or not to verify SSL certificates for S3 connection. (Default: None)
1099
+ This decorator is useful for handling transient errors, such as networking issues.
1100
+ If your task contains operations that can't be retried safely, e.g. database updates,
1101
+ it is advisable to annotate it with `@retry(times=0)`.
1102
+
1103
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
1104
+ decorator will execute a no-op task after all retries have been exhausted,
1105
+ ensuring that the flow execution can continue.
1106
+
1107
+ Parameters
1108
+ ----------
1109
+ times : int, default 3
1110
+ Number of times to retry this task.
1111
+ minutes_between_retries : int, default 2
1112
+ Number of minutes between retries.
1432
1113
  """
1433
1114
  ...
1434
1115
 
1435
- def project(*, name: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1116
+ @typing.overload
1117
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1118
+ ...
1119
+
1120
+ @typing.overload
1121
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1122
+ ...
1123
+
1124
+ 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):
1436
1125
  """
1437
- Specifies what flows belong to the same project.
1126
+ Specifies the number of times the task corresponding
1127
+ to a step needs to be retried.
1438
1128
 
1439
- A project-specific namespace is created for all flows that
1440
- use the same `@project(name)`.
1129
+ This decorator is useful for handling transient errors, such as networking issues.
1130
+ If your task contains operations that can't be retried safely, e.g. database updates,
1131
+ it is advisable to annotate it with `@retry(times=0)`.
1132
+
1133
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
1134
+ decorator will execute a no-op task after all retries have been exhausted,
1135
+ ensuring that the flow execution can continue.
1441
1136
 
1442
1137
  Parameters
1443
1138
  ----------
1444
- name : str
1445
- Project name. Make sure that the name is unique amongst all
1446
- projects that use the same production scheduler. The name may
1447
- contain only lowercase alphanumeric characters and underscores.
1139
+ times : int, default 3
1140
+ Number of times to retry this task.
1141
+ minutes_between_retries : int, default 2
1142
+ Number of minutes between retries.
1143
+ """
1144
+ ...
1145
+
1146
+ @typing.overload
1147
+ def secrets(*, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = []) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
1148
+ """
1149
+ Specifies secrets to be retrieved and injected as environment variables prior to
1150
+ the execution of a step.
1448
1151
 
1152
+ Parameters
1153
+ ----------
1154
+ sources : List[Union[str, Dict[str, Any]]], default: []
1155
+ List of secret specs, defining how the secrets are to be retrieved
1156
+ """
1157
+ ...
1158
+
1159
+ @typing.overload
1160
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1161
+ ...
1162
+
1163
+ @typing.overload
1164
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1165
+ ...
1166
+
1167
+ def secrets(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = []):
1168
+ """
1169
+ Specifies secrets to be retrieved and injected as environment variables prior to
1170
+ the execution of a step.
1449
1171
 
1172
+ Parameters
1173
+ ----------
1174
+ sources : List[Union[str, Dict[str, Any]]], default: []
1175
+ List of secret specs, defining how the secrets are to be retrieved
1450
1176
  """
1451
1177
  ...
1452
1178
 
1453
- 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]]:
1179
+ @typing.overload
1180
+ def resources(*, cpu: int = 1, gpu: int = 0, 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]]]:
1454
1181
  """
1455
- 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.
1456
- 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.
1182
+ Specifies the resources needed when executing this step.
1183
+
1184
+ Use `@resources` to specify the resource requirements
1185
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
1186
+
1187
+ You can choose the compute layer on the command line by executing e.g.
1188
+ ```
1189
+ python myflow.py run --with batch
1190
+ ```
1191
+ or
1192
+ ```
1193
+ python myflow.py run --with kubernetes
1194
+ ```
1195
+ which executes the flow on the desired system using the
1196
+ requirements specified in `@resources`.
1457
1197
 
1458
1198
  Parameters
1459
1199
  ----------
1460
- timeout : int
1461
- Time, in seconds before the task times out and fails. (Default: 3600)
1462
- poke_interval : int
1463
- Time in seconds that the job should wait in between each try. (Default: 60)
1464
- mode : str
1465
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1466
- exponential_backoff : bool
1467
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1468
- pool : str
1469
- the slot pool this task should run in,
1470
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1471
- soft_fail : bool
1472
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1473
- name : str
1474
- Name of the sensor on Airflow
1475
- description : str
1476
- Description of sensor in the Airflow UI
1477
- external_dag_id : str
1478
- The dag_id that contains the task you want to wait for.
1479
- external_task_ids : List[str]
1480
- The list of task_ids that you want to wait for.
1481
- If None (default value) the sensor waits for the DAG. (Default: None)
1482
- allowed_states : List[str]
1483
- Iterable of allowed states, (Default: ['success'])
1484
- failed_states : List[str]
1485
- Iterable of failed or dis-allowed states. (Default: None)
1486
- execution_delta : datetime.timedelta
1487
- time difference with the previous execution to look at,
1488
- the default is the same logical date as the current task or DAG. (Default: None)
1489
- check_existence: bool
1490
- Set to True to check if the external task exists or check if
1491
- the DAG to wait for exists. (Default: True)
1200
+ cpu : int, default 1
1201
+ Number of CPUs required for this step.
1202
+ gpu : int, default 0
1203
+ Number of GPUs required for this step.
1204
+ disk : int, optional, default None
1205
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
1206
+ memory : int, default 4096
1207
+ Memory size (in MB) required for this step.
1208
+ shared_memory : int, optional, default None
1209
+ The value for the size (in MiB) of the /dev/shm volume for this step.
1210
+ This parameter maps to the `--shm-size` option in Docker.
1492
1211
  """
1493
1212
  ...
1494
1213
 
1495
1214
  @typing.overload
1496
- 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]]:
1215
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1216
+ ...
1217
+
1218
+ @typing.overload
1219
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1220
+ ...
1221
+
1222
+ def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, cpu: int = 1, gpu: int = 0, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None):
1497
1223
  """
1498
- Specifies the Conda environment for all steps of the flow.
1224
+ Specifies the resources needed when executing this step.
1499
1225
 
1500
- Use `@conda_base` to set common libraries required by all
1501
- steps and use `@conda` to specify step-specific additions.
1226
+ Use `@resources` to specify the resource requirements
1227
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
1228
+
1229
+ You can choose the compute layer on the command line by executing e.g.
1230
+ ```
1231
+ python myflow.py run --with batch
1232
+ ```
1233
+ or
1234
+ ```
1235
+ python myflow.py run --with kubernetes
1236
+ ```
1237
+ which executes the flow on the desired system using the
1238
+ requirements specified in `@resources`.
1239
+
1240
+ Parameters
1241
+ ----------
1242
+ cpu : int, default 1
1243
+ Number of CPUs required for this step.
1244
+ gpu : int, default 0
1245
+ Number of GPUs required for this step.
1246
+ disk : int, optional, default None
1247
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
1248
+ memory : int, default 4096
1249
+ Memory size (in MB) required for this step.
1250
+ shared_memory : int, optional, default None
1251
+ The value for the size (in MiB) of the /dev/shm volume for this step.
1252
+ This parameter maps to the `--shm-size` option in Docker.
1253
+ """
1254
+ ...
1255
+
1256
+ @typing.overload
1257
+ 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]]]:
1258
+ """
1259
+ Specifies the Conda environment for the step.
1260
+
1261
+ Information in this decorator will augment any
1262
+ attributes set in the `@conda_base` flow-level decorator. Hence,
1263
+ you can use `@conda_base` to set packages required by all
1264
+ steps and use `@conda` to specify step-specific overrides.
1502
1265
 
1503
1266
  Parameters
1504
1267
  ----------
1505
1268
  packages : Dict[str, str], default {}
1506
- Packages to use for this flow. The key is the name of the package
1269
+ Packages to use for this step. The key is the name of the package
1507
1270
  and the value is the version to use.
1508
1271
  libraries : Dict[str, str], default {}
1509
1272
  Supported for backward compatibility. When used with packages, packages will take precedence.
@@ -1511,25 +1274,31 @@ def conda_base(*, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[s
1511
1274
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
1512
1275
  that the version used will correspond to the version of the Python interpreter used to start the run.
1513
1276
  disabled : bool, default False
1514
- If set to True, disables Conda.
1277
+ If set to True, disables @conda.
1515
1278
  """
1516
1279
  ...
1517
1280
 
1518
1281
  @typing.overload
1519
- def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1282
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1520
1283
  ...
1521
1284
 
1522
- 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):
1285
+ @typing.overload
1286
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1287
+ ...
1288
+
1289
+ 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):
1523
1290
  """
1524
- Specifies the Conda environment for all steps of the flow.
1291
+ Specifies the Conda environment for the step.
1525
1292
 
1526
- Use `@conda_base` to set common libraries required by all
1527
- steps and use `@conda` to specify step-specific additions.
1293
+ Information in this decorator will augment any
1294
+ attributes set in the `@conda_base` flow-level decorator. Hence,
1295
+ you can use `@conda_base` to set packages required by all
1296
+ steps and use `@conda` to specify step-specific overrides.
1528
1297
 
1529
1298
  Parameters
1530
1299
  ----------
1531
1300
  packages : Dict[str, str], default {}
1532
- Packages to use for this flow. The key is the name of the package
1301
+ Packages to use for this step. The key is the name of the package
1533
1302
  and the value is the version to use.
1534
1303
  libraries : Dict[str, str], default {}
1535
1304
  Supported for backward compatibility. When used with packages, packages will take precedence.
@@ -1537,21 +1306,55 @@ def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packa
1537
1306
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
1538
1307
  that the version used will correspond to the version of the Python interpreter used to start the run.
1539
1308
  disabled : bool, default False
1540
- If set to True, disables Conda.
1309
+ If set to True, disables @conda.
1541
1310
  """
1542
1311
  ...
1543
1312
 
1544
1313
  @typing.overload
1545
- def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1314
+ 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]]]:
1546
1315
  """
1547
- Specifies the PyPI packages for all steps of the flow.
1316
+ Specifies environment variables to be set prior to the execution of a step.
1548
1317
 
1549
- Use `@pypi_base` to set common packages required by all
1318
+ Parameters
1319
+ ----------
1320
+ vars : Dict[str, str], default {}
1321
+ Dictionary of environment variables to set.
1322
+ """
1323
+ ...
1324
+
1325
+ @typing.overload
1326
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1327
+ ...
1328
+
1329
+ @typing.overload
1330
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1331
+ ...
1332
+
1333
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
1334
+ """
1335
+ Specifies environment variables to be set prior to the execution of a step.
1336
+
1337
+ Parameters
1338
+ ----------
1339
+ vars : Dict[str, str], default {}
1340
+ Dictionary of environment variables to set.
1341
+ """
1342
+ ...
1343
+
1344
+ @typing.overload
1345
+ 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]]]:
1346
+ """
1347
+ Specifies the PyPI packages for the step.
1348
+
1349
+ Information in this decorator will augment any
1350
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
1351
+ you can use `@pypi_base` to set packages required by all
1550
1352
  steps and use `@pypi` to specify step-specific overrides.
1353
+
1551
1354
  Parameters
1552
1355
  ----------
1553
1356
  packages : Dict[str, str], default: {}
1554
- Packages to use for this flow. The key is the name of the package
1357
+ Packages to use for this step. The key is the name of the package
1555
1358
  and the value is the version to use.
1556
1359
  python : str, optional, default: None
1557
1360
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
@@ -1560,19 +1363,26 @@ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[s
1560
1363
  ...
1561
1364
 
1562
1365
  @typing.overload
1563
- def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1366
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1564
1367
  ...
1565
1368
 
1566
- def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1369
+ @typing.overload
1370
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1371
+ ...
1372
+
1373
+ 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):
1567
1374
  """
1568
- Specifies the PyPI packages for all steps of the flow.
1375
+ Specifies the PyPI packages for the step.
1569
1376
 
1570
- Use `@pypi_base` to set common packages required by all
1377
+ Information in this decorator will augment any
1378
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
1379
+ you can use `@pypi_base` to set packages required by all
1571
1380
  steps and use `@pypi` to specify step-specific overrides.
1381
+
1572
1382
  Parameters
1573
1383
  ----------
1574
1384
  packages : Dict[str, str], default: {}
1575
- Packages to use for this flow. The key is the name of the package
1385
+ Packages to use for this step. The key is the name of the package
1576
1386
  and the value is the version to use.
1577
1387
  python : str, optional, default: None
1578
1388
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
@@ -1580,98 +1390,133 @@ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packag
1580
1390
  """
1581
1391
  ...
1582
1392
 
1583
- @typing.overload
1584
- 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]]:
1393
+ 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]]:
1585
1394
  """
1586
- Specifies the event(s) that this flow depends on.
1587
-
1588
- ```
1589
- @trigger(event='foo')
1590
- ```
1591
- or
1592
- ```
1593
- @trigger(events=['foo', 'bar'])
1594
- ```
1395
+ 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.
1396
+ 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.
1595
1397
 
1596
- Additionally, you can specify the parameter mappings
1597
- to map event payload to Metaflow parameters for the flow.
1598
- ```
1599
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1600
- ```
1601
- or
1602
- ```
1603
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1604
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1605
- ```
1398
+ Parameters
1399
+ ----------
1400
+ timeout : int
1401
+ Time, in seconds before the task times out and fails. (Default: 3600)
1402
+ poke_interval : int
1403
+ Time in seconds that the job should wait in between each try. (Default: 60)
1404
+ mode : str
1405
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1406
+ exponential_backoff : bool
1407
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1408
+ pool : str
1409
+ the slot pool this task should run in,
1410
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1411
+ soft_fail : bool
1412
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1413
+ name : str
1414
+ Name of the sensor on Airflow
1415
+ description : str
1416
+ Description of sensor in the Airflow UI
1417
+ external_dag_id : str
1418
+ The dag_id that contains the task you want to wait for.
1419
+ external_task_ids : List[str]
1420
+ The list of task_ids that you want to wait for.
1421
+ If None (default value) the sensor waits for the DAG. (Default: None)
1422
+ allowed_states : List[str]
1423
+ Iterable of allowed states, (Default: ['success'])
1424
+ failed_states : List[str]
1425
+ Iterable of failed or dis-allowed states. (Default: None)
1426
+ execution_delta : datetime.timedelta
1427
+ time difference with the previous execution to look at,
1428
+ the default is the same logical date as the current task or DAG. (Default: None)
1429
+ check_existence: bool
1430
+ Set to True to check if the external task exists or check if
1431
+ the DAG to wait for exists. (Default: True)
1432
+ """
1433
+ ...
1434
+
1435
+ @typing.overload
1436
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1437
+ """
1438
+ Specifies the PyPI packages for all steps of the flow.
1606
1439
 
1607
- 'parameters' can also be a list of strings and tuples like so:
1608
- ```
1609
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1610
- ```
1611
- This is equivalent to:
1612
- ```
1613
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1614
- ```
1440
+ Use `@pypi_base` to set common packages required by all
1441
+ steps and use `@pypi` to specify step-specific overrides.
1442
+ Parameters
1443
+ ----------
1444
+ packages : Dict[str, str], default: {}
1445
+ Packages to use for this flow. The key is the name of the package
1446
+ and the value is the version to use.
1447
+ python : str, optional, default: None
1448
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1449
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1450
+ """
1451
+ ...
1452
+
1453
+ @typing.overload
1454
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1455
+ ...
1456
+
1457
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1458
+ """
1459
+ Specifies the PyPI packages for all steps of the flow.
1615
1460
 
1461
+ Use `@pypi_base` to set common packages required by all
1462
+ steps and use `@pypi` to specify step-specific overrides.
1616
1463
  Parameters
1617
1464
  ----------
1618
- event : Union[str, Dict[str, Any]], optional, default None
1619
- Event dependency for this flow.
1620
- events : List[Union[str, Dict[str, Any]]], default []
1621
- Events dependency for this flow.
1622
- options : Dict[str, Any], default {}
1623
- Backend-specific configuration for tuning eventing behavior.
1465
+ packages : Dict[str, str], default: {}
1466
+ Packages to use for this flow. The key is the name of the package
1467
+ and the value is the version to use.
1468
+ python : str, optional, default: None
1469
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1470
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1471
+ """
1472
+ ...
1473
+
1474
+ @typing.overload
1475
+ 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]]:
1476
+ """
1477
+ Specifies the Conda environment for all steps of the flow.
1624
1478
 
1479
+ Use `@conda_base` to set common libraries required by all
1480
+ steps and use `@conda` to specify step-specific additions.
1625
1481
 
1482
+ Parameters
1483
+ ----------
1484
+ packages : Dict[str, str], default {}
1485
+ Packages to use for this flow. The key is the name of the package
1486
+ and the value is the version to use.
1487
+ libraries : Dict[str, str], default {}
1488
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1489
+ python : str, optional, default None
1490
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1491
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1492
+ disabled : bool, default False
1493
+ If set to True, disables Conda.
1626
1494
  """
1627
1495
  ...
1628
1496
 
1629
1497
  @typing.overload
1630
- def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1498
+ def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1631
1499
  ...
1632
1500
 
1633
- 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] = {}):
1501
+ 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):
1634
1502
  """
1635
- Specifies the event(s) that this flow depends on.
1636
-
1637
- ```
1638
- @trigger(event='foo')
1639
- ```
1640
- or
1641
- ```
1642
- @trigger(events=['foo', 'bar'])
1643
- ```
1644
-
1645
- Additionally, you can specify the parameter mappings
1646
- to map event payload to Metaflow parameters for the flow.
1647
- ```
1648
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1649
- ```
1650
- or
1651
- ```
1652
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1653
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1654
- ```
1503
+ Specifies the Conda environment for all steps of the flow.
1655
1504
 
1656
- 'parameters' can also be a list of strings and tuples like so:
1657
- ```
1658
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1659
- ```
1660
- This is equivalent to:
1661
- ```
1662
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1663
- ```
1505
+ Use `@conda_base` to set common libraries required by all
1506
+ steps and use `@conda` to specify step-specific additions.
1664
1507
 
1665
1508
  Parameters
1666
1509
  ----------
1667
- event : Union[str, Dict[str, Any]], optional, default None
1668
- Event dependency for this flow.
1669
- events : List[Union[str, Dict[str, Any]]], default []
1670
- Events dependency for this flow.
1671
- options : Dict[str, Any], default {}
1672
- Backend-specific configuration for tuning eventing behavior.
1673
-
1674
-
1510
+ packages : Dict[str, str], default {}
1511
+ Packages to use for this flow. The key is the name of the package
1512
+ and the value is the version to use.
1513
+ libraries : Dict[str, str], default {}
1514
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1515
+ python : str, optional, default None
1516
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1517
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1518
+ disabled : bool, default False
1519
+ If set to True, disables Conda.
1675
1520
  """
1676
1521
  ...
1677
1522
 
@@ -1827,6 +1672,161 @@ def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *
1827
1672
  """
1828
1673
  ...
1829
1674
 
1675
+ @typing.overload
1676
+ 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]]:
1677
+ """
1678
+ Specifies the event(s) that this flow depends on.
1679
+
1680
+ ```
1681
+ @trigger(event='foo')
1682
+ ```
1683
+ or
1684
+ ```
1685
+ @trigger(events=['foo', 'bar'])
1686
+ ```
1687
+
1688
+ Additionally, you can specify the parameter mappings
1689
+ to map event payload to Metaflow parameters for the flow.
1690
+ ```
1691
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1692
+ ```
1693
+ or
1694
+ ```
1695
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1696
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1697
+ ```
1698
+
1699
+ 'parameters' can also be a list of strings and tuples like so:
1700
+ ```
1701
+ @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1702
+ ```
1703
+ This is equivalent to:
1704
+ ```
1705
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1706
+ ```
1707
+
1708
+ Parameters
1709
+ ----------
1710
+ event : Union[str, Dict[str, Any]], optional, default None
1711
+ Event dependency for this flow.
1712
+ events : List[Union[str, Dict[str, Any]]], default []
1713
+ Events dependency for this flow.
1714
+ options : Dict[str, Any], default {}
1715
+ Backend-specific configuration for tuning eventing behavior.
1716
+
1717
+
1718
+ """
1719
+ ...
1720
+
1721
+ @typing.overload
1722
+ def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1723
+ ...
1724
+
1725
+ 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] = {}):
1726
+ """
1727
+ Specifies the event(s) that this flow depends on.
1728
+
1729
+ ```
1730
+ @trigger(event='foo')
1731
+ ```
1732
+ or
1733
+ ```
1734
+ @trigger(events=['foo', 'bar'])
1735
+ ```
1736
+
1737
+ Additionally, you can specify the parameter mappings
1738
+ to map event payload to Metaflow parameters for the flow.
1739
+ ```
1740
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1741
+ ```
1742
+ or
1743
+ ```
1744
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1745
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1746
+ ```
1747
+
1748
+ 'parameters' can also be a list of strings and tuples like so:
1749
+ ```
1750
+ @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1751
+ ```
1752
+ This is equivalent to:
1753
+ ```
1754
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1755
+ ```
1756
+
1757
+ Parameters
1758
+ ----------
1759
+ event : Union[str, Dict[str, Any]], optional, default None
1760
+ Event dependency for this flow.
1761
+ events : List[Union[str, Dict[str, Any]]], default []
1762
+ Events dependency for this flow.
1763
+ options : Dict[str, Any], default {}
1764
+ Backend-specific configuration for tuning eventing behavior.
1765
+
1766
+
1767
+ """
1768
+ ...
1769
+
1770
+ def project(*, name: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1771
+ """
1772
+ Specifies what flows belong to the same project.
1773
+
1774
+ A project-specific namespace is created for all flows that
1775
+ use the same `@project(name)`.
1776
+
1777
+ Parameters
1778
+ ----------
1779
+ name : str
1780
+ Project name. Make sure that the name is unique amongst all
1781
+ projects that use the same production scheduler. The name may
1782
+ contain only lowercase alphanumeric characters and underscores.
1783
+
1784
+
1785
+ """
1786
+ ...
1787
+
1788
+ 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]]:
1789
+ """
1790
+ 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)
1791
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1792
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1793
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
1794
+ starts only after all sensors finish.
1795
+
1796
+ Parameters
1797
+ ----------
1798
+ timeout : int
1799
+ Time, in seconds before the task times out and fails. (Default: 3600)
1800
+ poke_interval : int
1801
+ Time in seconds that the job should wait in between each try. (Default: 60)
1802
+ mode : str
1803
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1804
+ exponential_backoff : bool
1805
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1806
+ pool : str
1807
+ the slot pool this task should run in,
1808
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1809
+ soft_fail : bool
1810
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1811
+ name : str
1812
+ Name of the sensor on Airflow
1813
+ description : str
1814
+ Description of sensor in the Airflow UI
1815
+ bucket_key : Union[str, List[str]]
1816
+ The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1817
+ When it's specified as a full s3:// url, please leave `bucket_name` as None
1818
+ bucket_name : str
1819
+ Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1820
+ When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1821
+ wildcard_match : bool
1822
+ whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1823
+ aws_conn_id : str
1824
+ a reference to the s3 connection on Airflow. (Default: None)
1825
+ verify : bool
1826
+ Whether or not to verify SSL certificates for S3 connection. (Default: None)
1827
+ """
1828
+ ...
1829
+
1830
1830
  def namespace(ns: typing.Optional[str]) -> typing.Optional[str]:
1831
1831
  """
1832
1832
  Switch namespace to the one provided.