metaflow-stubs 2.12.11__py2.py3-none-any.whl → 2.12.13__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 (150) hide show
  1. metaflow-stubs/__init__.pyi +542 -539
  2. metaflow-stubs/cards.pyi +4 -4
  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 +3 -3
  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 +6 -4
  16. metaflow-stubs/metaflow_current.pyi +3 -3
  17. metaflow-stubs/mflog/mflog.pyi +2 -2
  18. metaflow-stubs/multicore_utils.pyi +2 -2
  19. metaflow-stubs/parameters.pyi +3 -3
  20. metaflow-stubs/plugins/__init__.pyi +2 -2
  21. metaflow-stubs/plugins/airflow/__init__.pyi +2 -2
  22. metaflow-stubs/plugins/airflow/airflow.pyi +2 -2
  23. metaflow-stubs/plugins/airflow/airflow_cli.pyi +2 -2
  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 +2 -2
  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 +4 -4
  35. metaflow-stubs/plugins/argo/argo_workflows_cli.pyi +5 -5
  36. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +7 -7
  37. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +2 -2
  38. metaflow-stubs/plugins/aws/__init__.pyi +2 -2
  39. metaflow-stubs/plugins/aws/aws_client.pyi +2 -2
  40. metaflow-stubs/plugins/aws/aws_utils.pyi +5 -2
  41. metaflow-stubs/plugins/aws/batch/__init__.pyi +2 -2
  42. metaflow-stubs/plugins/aws/batch/batch.pyi +2 -2
  43. metaflow-stubs/plugins/aws/batch/batch_cli.pyi +2 -2
  44. metaflow-stubs/plugins/aws/batch/batch_client.pyi +2 -2
  45. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +2 -2
  46. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +2 -2
  47. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +4 -4
  48. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +2 -2
  49. metaflow-stubs/plugins/aws/step_functions/dynamo_db_client.pyi +2 -2
  50. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +2 -2
  51. metaflow-stubs/plugins/aws/step_functions/production_token.pyi +2 -2
  52. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +2 -2
  53. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +2 -2
  54. metaflow-stubs/plugins/aws/step_functions/step_functions_cli.pyi +3 -3
  55. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +2 -2
  56. metaflow-stubs/plugins/aws/step_functions/step_functions_decorator.pyi +2 -2
  57. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer.pyi +3 -3
  58. metaflow-stubs/plugins/azure/__init__.pyi +2 -2
  59. metaflow-stubs/plugins/azure/azure_credential.pyi +2 -2
  60. metaflow-stubs/plugins/azure/azure_exceptions.pyi +2 -2
  61. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +4 -4
  62. metaflow-stubs/plugins/azure/azure_utils.pyi +2 -2
  63. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +2 -2
  64. metaflow-stubs/plugins/azure/includefile_support.pyi +2 -2
  65. metaflow-stubs/plugins/cards/__init__.pyi +2 -2
  66. metaflow-stubs/plugins/cards/card_cli.pyi +3 -3
  67. metaflow-stubs/plugins/cards/card_client.pyi +4 -4
  68. metaflow-stubs/plugins/cards/card_creator.pyi +2 -2
  69. metaflow-stubs/plugins/cards/card_datastore.pyi +2 -2
  70. metaflow-stubs/plugins/cards/card_decorator.pyi +2 -2
  71. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +2 -2
  72. metaflow-stubs/plugins/cards/card_modules/basic.pyi +3 -3
  73. metaflow-stubs/plugins/cards/card_modules/card.pyi +2 -2
  74. metaflow-stubs/plugins/cards/card_modules/chevron/__init__.pyi +2 -2
  75. metaflow-stubs/plugins/cards/card_modules/chevron/main.pyi +2 -2
  76. metaflow-stubs/plugins/cards/card_modules/chevron/metadata.pyi +2 -2
  77. metaflow-stubs/plugins/cards/card_modules/chevron/renderer.pyi +2 -2
  78. metaflow-stubs/plugins/cards/card_modules/chevron/tokenizer.pyi +2 -2
  79. metaflow-stubs/plugins/cards/card_modules/components.pyi +3 -3
  80. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +2 -2
  81. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +2 -2
  82. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +3 -3
  83. metaflow-stubs/plugins/cards/card_resolver.pyi +2 -2
  84. metaflow-stubs/plugins/cards/component_serializer.pyi +3 -3
  85. metaflow-stubs/plugins/cards/exception.pyi +2 -2
  86. metaflow-stubs/plugins/catch_decorator.pyi +2 -2
  87. metaflow-stubs/plugins/datatools/__init__.pyi +3 -3
  88. metaflow-stubs/plugins/datatools/local.pyi +2 -2
  89. metaflow-stubs/plugins/datatools/s3/__init__.pyi +3 -3
  90. metaflow-stubs/plugins/datatools/s3/s3.pyi +4 -4
  91. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +2 -2
  92. metaflow-stubs/plugins/datatools/s3/s3util.pyi +2 -2
  93. metaflow-stubs/plugins/debug_logger.pyi +2 -2
  94. metaflow-stubs/plugins/debug_monitor.pyi +2 -2
  95. metaflow-stubs/plugins/environment_decorator.pyi +2 -2
  96. metaflow-stubs/plugins/events_decorator.pyi +2 -2
  97. metaflow-stubs/plugins/frameworks/__init__.pyi +2 -2
  98. metaflow-stubs/plugins/frameworks/pytorch.pyi +2 -2
  99. metaflow-stubs/plugins/gcp/__init__.pyi +2 -2
  100. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +4 -4
  101. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +2 -2
  102. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +2 -2
  103. metaflow-stubs/plugins/gcp/gs_utils.pyi +2 -2
  104. metaflow-stubs/plugins/gcp/includefile_support.pyi +2 -2
  105. metaflow-stubs/plugins/kubernetes/__init__.pyi +2 -2
  106. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +2 -2
  107. metaflow-stubs/plugins/kubernetes/kubernetes_cli.pyi +2 -2
  108. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +2 -2
  109. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +2 -2
  110. metaflow-stubs/plugins/kubernetes/kubernetes_job.pyi +2 -2
  111. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +2 -2
  112. metaflow-stubs/plugins/logs_cli.pyi +3 -3
  113. metaflow-stubs/plugins/package_cli.pyi +2 -2
  114. metaflow-stubs/plugins/parallel_decorator.pyi +2 -2
  115. metaflow-stubs/plugins/project_decorator.pyi +2 -2
  116. metaflow-stubs/plugins/pypi/__init__.pyi +2 -2
  117. metaflow-stubs/plugins/pypi/conda_decorator.pyi +2 -2
  118. metaflow-stubs/plugins/pypi/conda_environment.pyi +3 -3
  119. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +2 -2
  120. metaflow-stubs/plugins/pypi/pypi_environment.pyi +2 -2
  121. metaflow-stubs/plugins/pypi/utils.pyi +2 -2
  122. metaflow-stubs/plugins/resources_decorator.pyi +2 -2
  123. metaflow-stubs/plugins/retry_decorator.pyi +2 -2
  124. metaflow-stubs/plugins/secrets/__init__.pyi +2 -2
  125. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +3 -3
  126. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +2 -2
  127. metaflow-stubs/plugins/storage_executor.pyi +2 -2
  128. metaflow-stubs/plugins/tag_cli.pyi +3 -3
  129. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +2 -2
  130. metaflow-stubs/plugins/timeout_decorator.pyi +2 -2
  131. metaflow-stubs/procpoll.pyi +2 -2
  132. metaflow-stubs/pylint_wrapper.pyi +2 -2
  133. metaflow-stubs/runner/__init__.pyi +2 -2
  134. metaflow-stubs/runner/deployer.pyi +2 -2
  135. metaflow-stubs/runner/metaflow_runner.pyi +3 -3
  136. metaflow-stubs/runner/nbdeploy.pyi +2 -2
  137. metaflow-stubs/runner/nbrun.pyi +2 -2
  138. metaflow-stubs/runner/subprocess_manager.pyi +2 -2
  139. metaflow-stubs/runner/utils.pyi +2 -2
  140. metaflow-stubs/system/__init__.pyi +3 -3
  141. metaflow-stubs/system/system_logger.pyi +2 -2
  142. metaflow-stubs/system/system_monitor.pyi +3 -3
  143. metaflow-stubs/tagging_util.pyi +2 -2
  144. metaflow-stubs/tuple_util.pyi +2 -2
  145. metaflow-stubs/version.pyi +2 -2
  146. {metaflow_stubs-2.12.11.dist-info → metaflow_stubs-2.12.13.dist-info}/METADATA +2 -2
  147. metaflow_stubs-2.12.13.dist-info/RECORD +150 -0
  148. {metaflow_stubs-2.12.11.dist-info → metaflow_stubs-2.12.13.dist-info}/WHEEL +1 -1
  149. metaflow_stubs-2.12.11.dist-info/RECORD +0 -150
  150. {metaflow_stubs-2.12.11.dist-info → metaflow_stubs-2.12.13.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.11 #
4
- # Generated on 2024-08-09T01:05:47.258349 #
3
+ # MF version: 2.12.13 #
4
+ # Generated on 2024-08-21T05:55:53.002542 #
5
5
  ##################################################################################
6
6
 
7
7
  from __future__ import annotations
8
8
 
9
9
  import typing
10
10
  if typing.TYPE_CHECKING:
11
+ import typing
12
+ import metaflow.client.core
13
+ import metaflow.parameters
11
14
  import io
12
15
  import metaflow.metaflow_current
13
- import typing
14
16
  import metaflow.flowspec
15
- import metaflow.runner.metaflow_runner
16
- import metaflow.datastore.inputs
17
- import metaflow._vendor.click.types
17
+ import datetime
18
18
  import metaflow.events
19
- import metaflow.parameters
19
+ import metaflow.datastore.inputs
20
20
  import metaflow.plugins.datatools.s3.s3
21
- import datetime
22
- import metaflow.client.core
21
+ import metaflow.runner.metaflow_runner
22
+ import metaflow._vendor.click.types
23
23
  FlowSpecDerived = typing.TypeVar("FlowSpecDerived", bound="FlowSpec", contravariant=False, covariant=False)
24
24
  StepFlag = typing.NewType("StepFlag", bool)
25
25
 
@@ -728,33 +728,51 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
728
728
  ...
729
729
 
730
730
  @typing.overload
731
- 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]]]:
731
+ 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]]]:
732
732
  """
733
- Specifies environment variables to be set prior to the execution of a step.
733
+ Specifies the PyPI packages for the step.
734
+
735
+ Information in this decorator will augment any
736
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
737
+ you can use `@pypi_base` to set packages required by all
738
+ steps and use `@pypi` to specify step-specific overrides.
734
739
 
735
740
  Parameters
736
741
  ----------
737
- vars : Dict[str, str], default {}
738
- Dictionary of environment variables to set.
742
+ packages : Dict[str, str], default: {}
743
+ Packages to use for this step. The key is the name of the package
744
+ and the value is the version to use.
745
+ python : str, optional, default: None
746
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
747
+ that the version used will correspond to the version of the Python interpreter used to start the run.
739
748
  """
740
749
  ...
741
750
 
742
751
  @typing.overload
743
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
752
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
744
753
  ...
745
754
 
746
755
  @typing.overload
747
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
756
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
748
757
  ...
749
758
 
750
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
759
+ 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):
751
760
  """
752
- Specifies environment variables to be set prior to the execution of a step.
761
+ Specifies the PyPI packages for the step.
762
+
763
+ Information in this decorator will augment any
764
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
765
+ you can use `@pypi_base` to set packages required by all
766
+ steps and use `@pypi` to specify step-specific overrides.
753
767
 
754
768
  Parameters
755
769
  ----------
756
- vars : Dict[str, str], default {}
757
- Dictionary of environment variables to set.
770
+ packages : Dict[str, str], default: {}
771
+ Packages to use for this step. The key is the name of the package
772
+ and the value is the version to use.
773
+ python : str, optional, default: None
774
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
775
+ that the version used will correspond to the version of the Python interpreter used to start the run.
758
776
  """
759
777
  ...
760
778
 
@@ -815,6 +833,68 @@ def conda(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
815
833
  """
816
834
  ...
817
835
 
836
+ 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, compute_pool: typing.Optional[str] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
837
+ """
838
+ Specifies that this step should execute on Kubernetes.
839
+
840
+ Parameters
841
+ ----------
842
+ cpu : int, default 1
843
+ Number of CPUs required for this step. If `@resources` is
844
+ also present, the maximum value from all decorators is used.
845
+ memory : int, default 4096
846
+ Memory size (in MB) required for this step. If
847
+ `@resources` is also present, the maximum value from all decorators is
848
+ used.
849
+ disk : int, default 10240
850
+ Disk size (in MB) required for this step. If
851
+ `@resources` is also present, the maximum value from all decorators is
852
+ used.
853
+ image : str, optional, default None
854
+ Docker image to use when launching on Kubernetes. If not specified, and
855
+ METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
856
+ not, a default Docker image mapping to the current version of Python is used.
857
+ image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
858
+ If given, the imagePullPolicy to be applied to the Docker image of the step.
859
+ service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
860
+ Kubernetes service account to use when launching pod in Kubernetes.
861
+ secrets : List[str], optional, default None
862
+ Kubernetes secrets to use when launching pod in Kubernetes. These
863
+ secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
864
+ in Metaflow configuration.
865
+ namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
866
+ Kubernetes namespace to use when launching pod in Kubernetes.
867
+ gpu : int, optional, default None
868
+ Number of GPUs required for this step. A value of zero implies that
869
+ the scheduled node should not have GPUs.
870
+ gpu_vendor : str, default KUBERNETES_GPU_VENDOR
871
+ The vendor of the GPUs to be used for this step.
872
+ tolerations : List[str], default []
873
+ The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
874
+ Kubernetes tolerations to use when launching pod in Kubernetes.
875
+ use_tmpfs : bool, default False
876
+ This enables an explicit tmpfs mount for this step.
877
+ tmpfs_tempdir : bool, default True
878
+ sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
879
+ tmpfs_size : int, optional, default: None
880
+ The value for the size (in MiB) of the tmpfs mount for this step.
881
+ This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
882
+ memory allocated for this step.
883
+ tmpfs_path : str, optional, default /metaflow_temp
884
+ Path to tmpfs mount for this step.
885
+ persistent_volume_claims : Dict[str, str], optional, default None
886
+ A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
887
+ volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
888
+ shared_memory: int, optional
889
+ Shared memory size (in MiB) required for this step
890
+ port: int, optional
891
+ Port number to specify in the Kubernetes job object
892
+ compute_pool : str, optional, default None
893
+ Compute pool to be used for for this step.
894
+ If not specified, any accessible compute pool within the perimeter is used.
895
+ """
896
+ ...
897
+
818
898
  @typing.overload
819
899
  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]]]:
820
900
  """
@@ -869,255 +949,108 @@ def retry(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
869
949
  ...
870
950
 
871
951
  @typing.overload
872
- 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]]]:
952
+ 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]]]:
873
953
  """
874
- Specifies that this step should execute on [AWS Batch](https://aws.amazon.com/batch/).
954
+ Specifies a timeout for your step.
955
+
956
+ This decorator is useful if this step may hang indefinitely.
957
+
958
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
959
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
960
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
961
+
962
+ Note that all the values specified in parameters are added together so if you specify
963
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
964
+
965
+ Parameters
966
+ ----------
967
+ seconds : int, default 0
968
+ Number of seconds to wait prior to timing out.
969
+ minutes : int, default 0
970
+ Number of minutes to wait prior to timing out.
971
+ hours : int, default 0
972
+ Number of hours to wait prior to timing out.
973
+ """
974
+ ...
975
+
976
+ @typing.overload
977
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
978
+ ...
979
+
980
+ @typing.overload
981
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
982
+ ...
983
+
984
+ 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):
985
+ """
986
+ Specifies a timeout for your step.
987
+
988
+ This decorator is useful if this step may hang indefinitely.
989
+
990
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
991
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
992
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
993
+
994
+ Note that all the values specified in parameters are added together so if you specify
995
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
996
+
997
+ Parameters
998
+ ----------
999
+ seconds : int, default 0
1000
+ Number of seconds to wait prior to timing out.
1001
+ minutes : int, default 0
1002
+ Number of minutes to wait prior to timing out.
1003
+ hours : int, default 0
1004
+ Number of hours to wait prior to timing out.
1005
+ """
1006
+ ...
1007
+
1008
+ @typing.overload
1009
+ 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]]]:
1010
+ """
1011
+ Specifies the resources needed when executing this step.
1012
+
1013
+ Use `@resources` to specify the resource requirements
1014
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
1015
+
1016
+ You can choose the compute layer on the command line by executing e.g.
1017
+ ```
1018
+ python myflow.py run --with batch
1019
+ ```
1020
+ or
1021
+ ```
1022
+ python myflow.py run --with kubernetes
1023
+ ```
1024
+ which executes the flow on the desired system using the
1025
+ requirements specified in `@resources`.
875
1026
 
876
1027
  Parameters
877
1028
  ----------
878
1029
  cpu : int, default 1
879
- Number of CPUs required for this step. If `@resources` is
880
- also present, the maximum value from all decorators is used.
1030
+ Number of CPUs required for this step.
881
1031
  gpu : int, default 0
882
- Number of GPUs required for this step. If `@resources` is
883
- also present, the maximum value from all decorators is used.
1032
+ Number of GPUs required for this step.
1033
+ disk : int, optional, default None
1034
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
884
1035
  memory : int, default 4096
885
- Memory size (in MB) required for this step. If
886
- `@resources` is also present, the maximum value from all decorators is
887
- used.
888
- image : str, optional, default None
889
- Docker image to use when launching on AWS Batch. If not specified, and
890
- METAFLOW_BATCH_CONTAINER_IMAGE is specified, that image is used. If
891
- not, a default Docker image mapping to the current version of Python is used.
892
- queue : str, default METAFLOW_BATCH_JOB_QUEUE
893
- AWS Batch Job Queue to submit the job to.
894
- iam_role : str, default METAFLOW_ECS_S3_ACCESS_IAM_ROLE
895
- AWS IAM role that AWS Batch container uses to access AWS cloud resources.
896
- execution_role : str, default METAFLOW_ECS_FARGATE_EXECUTION_ROLE
897
- AWS IAM role that AWS Batch can use [to trigger AWS Fargate tasks]
898
- (https://docs.aws.amazon.com/batch/latest/userguide/execution-IAM-role.html).
1036
+ Memory size (in MB) required for this step.
899
1037
  shared_memory : int, optional, default None
900
1038
  The value for the size (in MiB) of the /dev/shm volume for this step.
901
1039
  This parameter maps to the `--shm-size` option in Docker.
902
- max_swap : int, optional, default None
903
- The total amount of swap memory (in MiB) a container can use for this
904
- step. This parameter is translated to the `--memory-swap` option in
905
- Docker where the value is the sum of the container memory plus the
906
- `max_swap` value.
907
- swappiness : int, optional, default None
908
- This allows you to tune memory swappiness behavior for this step.
909
- A swappiness value of 0 causes swapping not to happen unless absolutely
910
- necessary. A swappiness value of 100 causes pages to be swapped very
911
- aggressively. Accepted values are whole numbers between 0 and 100.
912
- use_tmpfs : bool, default False
913
- This enables an explicit tmpfs mount for this step. Note that tmpfs is
914
- not available on Fargate compute environments
915
- tmpfs_tempdir : bool, default True
916
- sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
917
- tmpfs_size : int, optional, default None
918
- The value for the size (in MiB) of the tmpfs mount for this step.
919
- This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
920
- memory allocated for this step.
921
- tmpfs_path : str, optional, default None
922
- Path to tmpfs mount for this step. Defaults to /metaflow_temp.
923
- inferentia : int, default 0
924
- Number of Inferentia chips required for this step.
925
- trainium : int, default None
926
- Alias for inferentia. Use only one of the two.
927
- efa : int, default 0
928
- Number of elastic fabric adapter network devices to attach to container
929
- ephemeral_storage : int, default None
930
- The total amount, in GiB, of ephemeral storage to set for the task, 21-200GiB.
931
- This is only relevant for Fargate compute environments
932
- log_driver: str, optional, default None
933
- The log driver to use for the Amazon ECS container.
934
- log_options: List[str], optional, default None
935
- List of strings containing options for the chosen log driver. The configurable values
936
- depend on the `log driver` chosen. Validation of these options is not supported yet.
937
- Example: [`awslogs-group:aws/batch/job`]
938
1040
  """
939
1041
  ...
940
1042
 
941
1043
  @typing.overload
942
- def batch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1044
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
943
1045
  ...
944
1046
 
945
1047
  @typing.overload
946
- def batch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1048
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
947
1049
  ...
948
1050
 
949
- def batch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, 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):
1051
+ 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):
950
1052
  """
951
- Specifies that this step should execute on [AWS Batch](https://aws.amazon.com/batch/).
952
-
953
- Parameters
954
- ----------
955
- cpu : int, default 1
956
- Number of CPUs required for this step. If `@resources` is
957
- also present, the maximum value from all decorators is used.
958
- gpu : int, default 0
959
- Number of GPUs required for this step. If `@resources` is
960
- also present, the maximum value from all decorators is used.
961
- memory : int, default 4096
962
- Memory size (in MB) required for this step. If
963
- `@resources` is also present, the maximum value from all decorators is
964
- used.
965
- image : str, optional, default None
966
- Docker image to use when launching on AWS Batch. If not specified, and
967
- METAFLOW_BATCH_CONTAINER_IMAGE is specified, that image is used. If
968
- not, a default Docker image mapping to the current version of Python is used.
969
- queue : str, default METAFLOW_BATCH_JOB_QUEUE
970
- AWS Batch Job Queue to submit the job to.
971
- iam_role : str, default METAFLOW_ECS_S3_ACCESS_IAM_ROLE
972
- AWS IAM role that AWS Batch container uses to access AWS cloud resources.
973
- execution_role : str, default METAFLOW_ECS_FARGATE_EXECUTION_ROLE
974
- AWS IAM role that AWS Batch can use [to trigger AWS Fargate tasks]
975
- (https://docs.aws.amazon.com/batch/latest/userguide/execution-IAM-role.html).
976
- shared_memory : int, optional, default None
977
- The value for the size (in MiB) of the /dev/shm volume for this step.
978
- This parameter maps to the `--shm-size` option in Docker.
979
- max_swap : int, optional, default None
980
- The total amount of swap memory (in MiB) a container can use for this
981
- step. This parameter is translated to the `--memory-swap` option in
982
- Docker where the value is the sum of the container memory plus the
983
- `max_swap` value.
984
- swappiness : int, optional, default None
985
- This allows you to tune memory swappiness behavior for this step.
986
- A swappiness value of 0 causes swapping not to happen unless absolutely
987
- necessary. A swappiness value of 100 causes pages to be swapped very
988
- aggressively. Accepted values are whole numbers between 0 and 100.
989
- use_tmpfs : bool, default False
990
- This enables an explicit tmpfs mount for this step. Note that tmpfs is
991
- not available on Fargate compute environments
992
- tmpfs_tempdir : bool, default True
993
- sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
994
- tmpfs_size : int, optional, default None
995
- The value for the size (in MiB) of the tmpfs mount for this step.
996
- This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
997
- memory allocated for this step.
998
- tmpfs_path : str, optional, default None
999
- Path to tmpfs mount for this step. Defaults to /metaflow_temp.
1000
- inferentia : int, default 0
1001
- Number of Inferentia chips required for this step.
1002
- trainium : int, default None
1003
- Alias for inferentia. Use only one of the two.
1004
- efa : int, default 0
1005
- Number of elastic fabric adapter network devices to attach to container
1006
- ephemeral_storage : int, default None
1007
- The total amount, in GiB, of ephemeral storage to set for the task, 21-200GiB.
1008
- This is only relevant for Fargate compute environments
1009
- log_driver: str, optional, default None
1010
- The log driver to use for the Amazon ECS container.
1011
- log_options: List[str], optional, default None
1012
- List of strings containing options for the chosen log driver. The configurable values
1013
- depend on the `log driver` chosen. Validation of these options is not supported yet.
1014
- Example: [`awslogs-group:aws/batch/job`]
1015
- """
1016
- ...
1017
-
1018
- @typing.overload
1019
- 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]]]:
1020
- """
1021
- Specifies a timeout for your step.
1022
-
1023
- This decorator is useful if this step may hang indefinitely.
1024
-
1025
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
1026
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
1027
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
1028
-
1029
- Note that all the values specified in parameters are added together so if you specify
1030
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
1031
-
1032
- Parameters
1033
- ----------
1034
- seconds : int, default 0
1035
- Number of seconds to wait prior to timing out.
1036
- minutes : int, default 0
1037
- Number of minutes to wait prior to timing out.
1038
- hours : int, default 0
1039
- Number of hours to wait prior to timing out.
1040
- """
1041
- ...
1042
-
1043
- @typing.overload
1044
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1045
- ...
1046
-
1047
- @typing.overload
1048
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1049
- ...
1050
-
1051
- 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):
1052
- """
1053
- Specifies a timeout for your step.
1054
-
1055
- This decorator is useful if this step may hang indefinitely.
1056
-
1057
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
1058
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
1059
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
1060
-
1061
- Note that all the values specified in parameters are added together so if you specify
1062
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
1063
-
1064
- Parameters
1065
- ----------
1066
- seconds : int, default 0
1067
- Number of seconds to wait prior to timing out.
1068
- minutes : int, default 0
1069
- Number of minutes to wait prior to timing out.
1070
- hours : int, default 0
1071
- Number of hours to wait prior to timing out.
1072
- """
1073
- ...
1074
-
1075
- @typing.overload
1076
- 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]]]:
1077
- """
1078
- Specifies the resources needed when executing this step.
1079
-
1080
- Use `@resources` to specify the resource requirements
1081
- independently of the specific compute layer (`@batch`, `@kubernetes`).
1082
-
1083
- You can choose the compute layer on the command line by executing e.g.
1084
- ```
1085
- python myflow.py run --with batch
1086
- ```
1087
- or
1088
- ```
1089
- python myflow.py run --with kubernetes
1090
- ```
1091
- which executes the flow on the desired system using the
1092
- requirements specified in `@resources`.
1093
-
1094
- Parameters
1095
- ----------
1096
- cpu : int, default 1
1097
- Number of CPUs required for this step.
1098
- gpu : int, default 0
1099
- Number of GPUs required for this step.
1100
- disk : int, optional, default None
1101
- Disk size (in MB) required for this step. Only applies on Kubernetes.
1102
- memory : int, default 4096
1103
- Memory size (in MB) required for this step.
1104
- shared_memory : int, optional, default None
1105
- The value for the size (in MiB) of the /dev/shm volume for this step.
1106
- This parameter maps to the `--shm-size` option in Docker.
1107
- """
1108
- ...
1109
-
1110
- @typing.overload
1111
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1112
- ...
1113
-
1114
- @typing.overload
1115
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1116
- ...
1117
-
1118
- 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):
1119
- """
1120
- Specifies the resources needed when executing this step.
1053
+ Specifies the resources needed when executing this step.
1121
1054
 
1122
1055
  Use `@resources` to specify the resource requirements
1123
1056
  independently of the specific compute layer (`@batch`, `@kubernetes`).
@@ -1150,51 +1083,35 @@ def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None]
1150
1083
  ...
1151
1084
 
1152
1085
  @typing.overload
1153
- 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]]]:
1086
+ 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]]]:
1154
1087
  """
1155
- Specifies that the step will success under all circumstances.
1156
-
1157
- The decorator will create an optional artifact, specified by `var`, which
1158
- contains the exception raised. You can use it to detect the presence
1159
- of errors, indicating that all happy-path artifacts produced by the step
1160
- are missing.
1088
+ Specifies secrets to be retrieved and injected as environment variables prior to
1089
+ the execution of a step.
1161
1090
 
1162
1091
  Parameters
1163
1092
  ----------
1164
- var : str, optional, default None
1165
- Name of the artifact in which to store the caught exception.
1166
- If not specified, the exception is not stored.
1167
- print_exception : bool, default True
1168
- Determines whether or not the exception is printed to
1169
- stdout when caught.
1093
+ sources : List[Union[str, Dict[str, Any]]], default: []
1094
+ List of secret specs, defining how the secrets are to be retrieved
1170
1095
  """
1171
1096
  ...
1172
1097
 
1173
1098
  @typing.overload
1174
- def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1099
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1175
1100
  ...
1176
1101
 
1177
1102
  @typing.overload
1178
- def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1103
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1179
1104
  ...
1180
1105
 
1181
- def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, var: typing.Optional[str] = None, print_exception: bool = True):
1106
+ 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]]] = []):
1182
1107
  """
1183
- Specifies that the step will success under all circumstances.
1184
-
1185
- The decorator will create an optional artifact, specified by `var`, which
1186
- contains the exception raised. You can use it to detect the presence
1187
- of errors, indicating that all happy-path artifacts produced by the step
1188
- are missing.
1108
+ Specifies secrets to be retrieved and injected as environment variables prior to
1109
+ the execution of a step.
1189
1110
 
1190
1111
  Parameters
1191
1112
  ----------
1192
- var : str, optional, default None
1193
- Name of the artifact in which to store the caught exception.
1194
- If not specified, the exception is not stored.
1195
- print_exception : bool, default True
1196
- Determines whether or not the exception is printed to
1197
- stdout when caught.
1113
+ sources : List[Union[str, Dict[str, Any]]], default: []
1114
+ List of secret specs, defining how the secrets are to be retrieved
1198
1115
  """
1199
1116
  ...
1200
1117
 
@@ -1249,95 +1166,181 @@ def card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typ
1249
1166
  """
1250
1167
  ...
1251
1168
 
1252
- 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]]]:
1169
+ @typing.overload
1170
+ 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]]]:
1253
1171
  """
1254
- Specifies that this step should execute on Kubernetes.
1172
+ Specifies environment variables to be set prior to the execution of a step.
1173
+
1174
+ Parameters
1175
+ ----------
1176
+ vars : Dict[str, str], default {}
1177
+ Dictionary of environment variables to set.
1178
+ """
1179
+ ...
1180
+
1181
+ @typing.overload
1182
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1183
+ ...
1184
+
1185
+ @typing.overload
1186
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1187
+ ...
1188
+
1189
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
1190
+ """
1191
+ Specifies environment variables to be set prior to the execution of a step.
1192
+
1193
+ Parameters
1194
+ ----------
1195
+ vars : Dict[str, str], default {}
1196
+ Dictionary of environment variables to set.
1197
+ """
1198
+ ...
1199
+
1200
+ @typing.overload
1201
+ 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]]]:
1202
+ """
1203
+ Specifies that this step should execute on [AWS Batch](https://aws.amazon.com/batch/).
1255
1204
 
1256
1205
  Parameters
1257
1206
  ----------
1258
1207
  cpu : int, default 1
1259
1208
  Number of CPUs required for this step. If `@resources` is
1260
1209
  also present, the maximum value from all decorators is used.
1210
+ gpu : int, default 0
1211
+ Number of GPUs required for this step. If `@resources` is
1212
+ also present, the maximum value from all decorators is used.
1261
1213
  memory : int, default 4096
1262
1214
  Memory size (in MB) required for this step. If
1263
1215
  `@resources` is also present, the maximum value from all decorators is
1264
1216
  used.
1265
- disk : int, default 10240
1266
- Disk size (in MB) required for this step. If
1267
- `@resources` is also present, the maximum value from all decorators is
1268
- used.
1269
1217
  image : str, optional, default None
1270
- Docker image to use when launching on Kubernetes. If not specified, and
1271
- METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
1218
+ Docker image to use when launching on AWS Batch. If not specified, and
1219
+ METAFLOW_BATCH_CONTAINER_IMAGE is specified, that image is used. If
1272
1220
  not, a default Docker image mapping to the current version of Python is used.
1273
- image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
1274
- If given, the imagePullPolicy to be applied to the Docker image of the step.
1275
- service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
1276
- Kubernetes service account to use when launching pod in Kubernetes.
1277
- secrets : List[str], optional, default None
1278
- Kubernetes secrets to use when launching pod in Kubernetes. These
1279
- secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
1280
- in Metaflow configuration.
1281
- namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
1282
- Kubernetes namespace to use when launching pod in Kubernetes.
1283
- gpu : int, optional, default None
1284
- Number of GPUs required for this step. A value of zero implies that
1285
- the scheduled node should not have GPUs.
1286
- gpu_vendor : str, default KUBERNETES_GPU_VENDOR
1287
- The vendor of the GPUs to be used for this step.
1288
- tolerations : List[str], default []
1289
- The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
1290
- Kubernetes tolerations to use when launching pod in Kubernetes.
1221
+ queue : str, default METAFLOW_BATCH_JOB_QUEUE
1222
+ AWS Batch Job Queue to submit the job to.
1223
+ iam_role : str, default METAFLOW_ECS_S3_ACCESS_IAM_ROLE
1224
+ AWS IAM role that AWS Batch container uses to access AWS cloud resources.
1225
+ execution_role : str, default METAFLOW_ECS_FARGATE_EXECUTION_ROLE
1226
+ AWS IAM role that AWS Batch can use [to trigger AWS Fargate tasks]
1227
+ (https://docs.aws.amazon.com/batch/latest/userguide/execution-IAM-role.html).
1228
+ shared_memory : int, optional, default None
1229
+ The value for the size (in MiB) of the /dev/shm volume for this step.
1230
+ This parameter maps to the `--shm-size` option in Docker.
1231
+ max_swap : int, optional, default None
1232
+ The total amount of swap memory (in MiB) a container can use for this
1233
+ step. This parameter is translated to the `--memory-swap` option in
1234
+ Docker where the value is the sum of the container memory plus the
1235
+ `max_swap` value.
1236
+ swappiness : int, optional, default None
1237
+ This allows you to tune memory swappiness behavior for this step.
1238
+ A swappiness value of 0 causes swapping not to happen unless absolutely
1239
+ necessary. A swappiness value of 100 causes pages to be swapped very
1240
+ aggressively. Accepted values are whole numbers between 0 and 100.
1291
1241
  use_tmpfs : bool, default False
1292
- This enables an explicit tmpfs mount for this step.
1242
+ This enables an explicit tmpfs mount for this step. Note that tmpfs is
1243
+ not available on Fargate compute environments
1293
1244
  tmpfs_tempdir : bool, default True
1294
1245
  sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
1295
- tmpfs_size : int, optional, default: None
1246
+ tmpfs_size : int, optional, default None
1296
1247
  The value for the size (in MiB) of the tmpfs mount for this step.
1297
1248
  This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
1298
1249
  memory allocated for this step.
1299
- tmpfs_path : str, optional, default /metaflow_temp
1300
- Path to tmpfs mount for this step.
1301
- persistent_volume_claims : Dict[str, str], optional, default None
1302
- A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
1303
- volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
1304
- shared_memory: int, optional
1305
- Shared memory size (in MiB) required for this step
1306
- port: int, optional
1307
- Port number to specify in the Kubernetes job object
1308
- """
1309
- ...
1310
-
1311
- @typing.overload
1312
- 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]]]:
1313
- """
1314
- Specifies secrets to be retrieved and injected as environment variables prior to
1315
- the execution of a step.
1316
-
1317
- Parameters
1318
- ----------
1319
- sources : List[Union[str, Dict[str, Any]]], default: []
1320
- List of secret specs, defining how the secrets are to be retrieved
1250
+ tmpfs_path : str, optional, default None
1251
+ Path to tmpfs mount for this step. Defaults to /metaflow_temp.
1252
+ inferentia : int, default 0
1253
+ Number of Inferentia chips required for this step.
1254
+ trainium : int, default None
1255
+ Alias for inferentia. Use only one of the two.
1256
+ efa : int, default 0
1257
+ Number of elastic fabric adapter network devices to attach to container
1258
+ ephemeral_storage : int, default None
1259
+ The total amount, in GiB, of ephemeral storage to set for the task, 21-200GiB.
1260
+ This is only relevant for Fargate compute environments
1261
+ log_driver: str, optional, default None
1262
+ The log driver to use for the Amazon ECS container.
1263
+ log_options: List[str], optional, default None
1264
+ List of strings containing options for the chosen log driver. The configurable values
1265
+ depend on the `log driver` chosen. Validation of these options is not supported yet.
1266
+ Example: [`awslogs-group:aws/batch/job`]
1321
1267
  """
1322
1268
  ...
1323
1269
 
1324
1270
  @typing.overload
1325
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1271
+ def batch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1326
1272
  ...
1327
1273
 
1328
1274
  @typing.overload
1329
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1275
+ def batch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1330
1276
  ...
1331
1277
 
1332
- 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]]] = []):
1278
+ def batch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, 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):
1333
1279
  """
1334
- Specifies secrets to be retrieved and injected as environment variables prior to
1335
- the execution of a step.
1280
+ Specifies that this step should execute on [AWS Batch](https://aws.amazon.com/batch/).
1336
1281
 
1337
1282
  Parameters
1338
1283
  ----------
1339
- sources : List[Union[str, Dict[str, Any]]], default: []
1340
- List of secret specs, defining how the secrets are to be retrieved
1284
+ cpu : int, default 1
1285
+ Number of CPUs required for this step. If `@resources` is
1286
+ also present, the maximum value from all decorators is used.
1287
+ gpu : int, default 0
1288
+ Number of GPUs required for this step. If `@resources` is
1289
+ also present, the maximum value from all decorators is used.
1290
+ memory : int, default 4096
1291
+ Memory size (in MB) required for this step. If
1292
+ `@resources` is also present, the maximum value from all decorators is
1293
+ used.
1294
+ image : str, optional, default None
1295
+ Docker image to use when launching on AWS Batch. If not specified, and
1296
+ METAFLOW_BATCH_CONTAINER_IMAGE is specified, that image is used. If
1297
+ not, a default Docker image mapping to the current version of Python is used.
1298
+ queue : str, default METAFLOW_BATCH_JOB_QUEUE
1299
+ AWS Batch Job Queue to submit the job to.
1300
+ iam_role : str, default METAFLOW_ECS_S3_ACCESS_IAM_ROLE
1301
+ AWS IAM role that AWS Batch container uses to access AWS cloud resources.
1302
+ execution_role : str, default METAFLOW_ECS_FARGATE_EXECUTION_ROLE
1303
+ AWS IAM role that AWS Batch can use [to trigger AWS Fargate tasks]
1304
+ (https://docs.aws.amazon.com/batch/latest/userguide/execution-IAM-role.html).
1305
+ shared_memory : int, optional, default None
1306
+ The value for the size (in MiB) of the /dev/shm volume for this step.
1307
+ This parameter maps to the `--shm-size` option in Docker.
1308
+ max_swap : int, optional, default None
1309
+ The total amount of swap memory (in MiB) a container can use for this
1310
+ step. This parameter is translated to the `--memory-swap` option in
1311
+ Docker where the value is the sum of the container memory plus the
1312
+ `max_swap` value.
1313
+ swappiness : int, optional, default None
1314
+ This allows you to tune memory swappiness behavior for this step.
1315
+ A swappiness value of 0 causes swapping not to happen unless absolutely
1316
+ necessary. A swappiness value of 100 causes pages to be swapped very
1317
+ aggressively. Accepted values are whole numbers between 0 and 100.
1318
+ use_tmpfs : bool, default False
1319
+ This enables an explicit tmpfs mount for this step. Note that tmpfs is
1320
+ not available on Fargate compute environments
1321
+ tmpfs_tempdir : bool, default True
1322
+ sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
1323
+ tmpfs_size : int, optional, default None
1324
+ The value for the size (in MiB) of the tmpfs mount for this step.
1325
+ This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
1326
+ memory allocated for this step.
1327
+ tmpfs_path : str, optional, default None
1328
+ Path to tmpfs mount for this step. Defaults to /metaflow_temp.
1329
+ inferentia : int, default 0
1330
+ Number of Inferentia chips required for this step.
1331
+ trainium : int, default None
1332
+ Alias for inferentia. Use only one of the two.
1333
+ efa : int, default 0
1334
+ Number of elastic fabric adapter network devices to attach to container
1335
+ ephemeral_storage : int, default None
1336
+ The total amount, in GiB, of ephemeral storage to set for the task, 21-200GiB.
1337
+ This is only relevant for Fargate compute environments
1338
+ log_driver: str, optional, default None
1339
+ The log driver to use for the Amazon ECS container.
1340
+ log_options: List[str], optional, default None
1341
+ List of strings containing options for the chosen log driver. The configurable values
1342
+ depend on the `log driver` chosen. Validation of these options is not supported yet.
1343
+ Example: [`awslogs-group:aws/batch/job`]
1341
1344
  """
1342
1345
  ...
1343
1346
 
@@ -1361,233 +1364,90 @@ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None],
1361
1364
  ...
1362
1365
 
1363
1366
  @typing.overload
1364
- 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]]]:
1365
- """
1366
- Specifies the PyPI packages for the step.
1367
-
1368
- Information in this decorator will augment any
1369
- attributes set in the `@pyi_base` flow-level decorator. Hence,
1370
- you can use `@pypi_base` to set packages required by all
1371
- steps and use `@pypi` to specify step-specific overrides.
1372
-
1373
- Parameters
1374
- ----------
1375
- packages : Dict[str, str], default: {}
1376
- Packages to use for this step. The key is the name of the package
1377
- and the value is the version to use.
1378
- python : str, optional, default: None
1379
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1380
- that the version used will correspond to the version of the Python interpreter used to start the run.
1381
- """
1382
- ...
1383
-
1384
- @typing.overload
1385
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1386
- ...
1387
-
1388
- @typing.overload
1389
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1390
- ...
1391
-
1392
- 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):
1393
- """
1394
- Specifies the PyPI packages for the step.
1395
-
1396
- Information in this decorator will augment any
1397
- attributes set in the `@pyi_base` flow-level decorator. Hence,
1398
- you can use `@pypi_base` to set packages required by all
1399
- steps and use `@pypi` to specify step-specific overrides.
1400
-
1401
- Parameters
1402
- ----------
1403
- packages : Dict[str, str], default: {}
1404
- Packages to use for this step. The key is the name of the package
1405
- and the value is the version to use.
1406
- python : str, optional, default: None
1407
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1408
- that the version used will correspond to the version of the Python interpreter used to start the run.
1409
- """
1410
- ...
1411
-
1412
- 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]]:
1367
+ 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]]]:
1413
1368
  """
1414
- 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)
1415
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1416
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1417
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
1418
- starts only after all sensors finish.
1369
+ Specifies that the step will success under all circumstances.
1419
1370
 
1420
- Parameters
1421
- ----------
1422
- timeout : int
1423
- Time, in seconds before the task times out and fails. (Default: 3600)
1424
- poke_interval : int
1425
- Time in seconds that the job should wait in between each try. (Default: 60)
1426
- mode : str
1427
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1428
- exponential_backoff : bool
1429
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1430
- pool : str
1431
- the slot pool this task should run in,
1432
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1433
- soft_fail : bool
1434
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1435
- name : str
1436
- Name of the sensor on Airflow
1437
- description : str
1438
- Description of sensor in the Airflow UI
1439
- bucket_key : Union[str, List[str]]
1440
- The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1441
- When it's specified as a full s3:// url, please leave `bucket_name` as None
1442
- bucket_name : str
1443
- Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1444
- When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1445
- wildcard_match : bool
1446
- whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1447
- aws_conn_id : str
1448
- a reference to the s3 connection on Airflow. (Default: None)
1449
- verify : bool
1450
- Whether or not to verify SSL certificates for S3 connection. (Default: None)
1451
- """
1452
- ...
1453
-
1454
- 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]]:
1455
- """
1456
- 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.
1457
- 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.
1371
+ The decorator will create an optional artifact, specified by `var`, which
1372
+ contains the exception raised. You can use it to detect the presence
1373
+ of errors, indicating that all happy-path artifacts produced by the step
1374
+ are missing.
1458
1375
 
1459
1376
  Parameters
1460
1377
  ----------
1461
- timeout : int
1462
- Time, in seconds before the task times out and fails. (Default: 3600)
1463
- poke_interval : int
1464
- Time in seconds that the job should wait in between each try. (Default: 60)
1465
- mode : str
1466
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1467
- exponential_backoff : bool
1468
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1469
- pool : str
1470
- the slot pool this task should run in,
1471
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1472
- soft_fail : bool
1473
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1474
- name : str
1475
- Name of the sensor on Airflow
1476
- description : str
1477
- Description of sensor in the Airflow UI
1478
- external_dag_id : str
1479
- The dag_id that contains the task you want to wait for.
1480
- external_task_ids : List[str]
1481
- The list of task_ids that you want to wait for.
1482
- If None (default value) the sensor waits for the DAG. (Default: None)
1483
- allowed_states : List[str]
1484
- Iterable of allowed states, (Default: ['success'])
1485
- failed_states : List[str]
1486
- Iterable of failed or dis-allowed states. (Default: None)
1487
- execution_delta : datetime.timedelta
1488
- time difference with the previous execution to look at,
1489
- the default is the same logical date as the current task or DAG. (Default: None)
1490
- check_existence: bool
1491
- Set to True to check if the external task exists or check if
1492
- the DAG to wait for exists. (Default: True)
1378
+ var : str, optional, default None
1379
+ Name of the artifact in which to store the caught exception.
1380
+ If not specified, the exception is not stored.
1381
+ print_exception : bool, default True
1382
+ Determines whether or not the exception is printed to
1383
+ stdout when caught.
1493
1384
  """
1494
1385
  ...
1495
1386
 
1496
1387
  @typing.overload
1497
- def schedule(*, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1498
- """
1499
- Specifies the times when the flow should be run when running on a
1500
- production scheduler.
1501
-
1502
- Parameters
1503
- ----------
1504
- hourly : bool, default False
1505
- Run the workflow hourly.
1506
- daily : bool, default True
1507
- Run the workflow daily.
1508
- weekly : bool, default False
1509
- Run the workflow weekly.
1510
- cron : str, optional, default None
1511
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1512
- specified by this expression.
1513
- timezone : str, optional, default None
1514
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1515
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1516
- """
1388
+ def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1517
1389
  ...
1518
1390
 
1519
1391
  @typing.overload
1520
- def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1392
+ def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1521
1393
  ...
1522
1394
 
1523
- def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None):
1395
+ def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, var: typing.Optional[str] = None, print_exception: bool = True):
1524
1396
  """
1525
- Specifies the times when the flow should be run when running on a
1526
- production scheduler.
1397
+ Specifies that the step will success under all circumstances.
1398
+
1399
+ The decorator will create an optional artifact, specified by `var`, which
1400
+ contains the exception raised. You can use it to detect the presence
1401
+ of errors, indicating that all happy-path artifacts produced by the step
1402
+ are missing.
1527
1403
 
1528
1404
  Parameters
1529
1405
  ----------
1530
- hourly : bool, default False
1531
- Run the workflow hourly.
1532
- daily : bool, default True
1533
- Run the workflow daily.
1534
- weekly : bool, default False
1535
- Run the workflow weekly.
1536
- cron : str, optional, default None
1537
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1538
- specified by this expression.
1539
- timezone : str, optional, default None
1540
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1541
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1406
+ var : str, optional, default None
1407
+ Name of the artifact in which to store the caught exception.
1408
+ If not specified, the exception is not stored.
1409
+ print_exception : bool, default True
1410
+ Determines whether or not the exception is printed to
1411
+ stdout when caught.
1542
1412
  """
1543
1413
  ...
1544
1414
 
1545
1415
  @typing.overload
1546
- 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]]:
1416
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1547
1417
  """
1548
- Specifies the Conda environment for all steps of the flow.
1549
-
1550
- Use `@conda_base` to set common libraries required by all
1551
- steps and use `@conda` to specify step-specific additions.
1418
+ Specifies the PyPI packages for all steps of the flow.
1552
1419
 
1420
+ Use `@pypi_base` to set common packages required by all
1421
+ steps and use `@pypi` to specify step-specific overrides.
1553
1422
  Parameters
1554
1423
  ----------
1555
- packages : Dict[str, str], default {}
1424
+ packages : Dict[str, str], default: {}
1556
1425
  Packages to use for this flow. The key is the name of the package
1557
1426
  and the value is the version to use.
1558
- libraries : Dict[str, str], default {}
1559
- Supported for backward compatibility. When used with packages, packages will take precedence.
1560
- python : str, optional, default None
1427
+ python : str, optional, default: None
1561
1428
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
1562
1429
  that the version used will correspond to the version of the Python interpreter used to start the run.
1563
- disabled : bool, default False
1564
- If set to True, disables Conda.
1565
1430
  """
1566
1431
  ...
1567
1432
 
1568
1433
  @typing.overload
1569
- def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1434
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1570
1435
  ...
1571
1436
 
1572
- 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):
1437
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1573
1438
  """
1574
- Specifies the Conda environment for all steps of the flow.
1575
-
1576
- Use `@conda_base` to set common libraries required by all
1577
- steps and use `@conda` to specify step-specific additions.
1439
+ Specifies the PyPI packages for all steps of the flow.
1578
1440
 
1441
+ Use `@pypi_base` to set common packages required by all
1442
+ steps and use `@pypi` to specify step-specific overrides.
1579
1443
  Parameters
1580
1444
  ----------
1581
- packages : Dict[str, str], default {}
1445
+ packages : Dict[str, str], default: {}
1582
1446
  Packages to use for this flow. The key is the name of the package
1583
1447
  and the value is the version to use.
1584
- libraries : Dict[str, str], default {}
1585
- Supported for backward compatibility. When used with packages, packages will take precedence.
1586
- python : str, optional, default None
1448
+ python : str, optional, default: None
1587
1449
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
1588
1450
  that the version used will correspond to the version of the Python interpreter used to start the run.
1589
- disabled : bool, default False
1590
- If set to True, disables Conda.
1591
1451
  """
1592
1452
  ...
1593
1453
 
@@ -1695,41 +1555,160 @@ def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *
1695
1555
  ...
1696
1556
 
1697
1557
  @typing.overload
1698
- def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1558
+ def schedule(*, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1699
1559
  """
1700
- Specifies the PyPI packages for all steps of the flow.
1560
+ Specifies the times when the flow should be run when running on a
1561
+ production scheduler.
1701
1562
 
1702
- Use `@pypi_base` to set common packages required by all
1703
- steps and use `@pypi` to specify step-specific overrides.
1704
1563
  Parameters
1705
1564
  ----------
1706
- packages : Dict[str, str], default: {}
1565
+ hourly : bool, default False
1566
+ Run the workflow hourly.
1567
+ daily : bool, default True
1568
+ Run the workflow daily.
1569
+ weekly : bool, default False
1570
+ Run the workflow weekly.
1571
+ cron : str, optional, default None
1572
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1573
+ specified by this expression.
1574
+ timezone : str, optional, default None
1575
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1576
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1577
+ """
1578
+ ...
1579
+
1580
+ @typing.overload
1581
+ def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1582
+ ...
1583
+
1584
+ def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None):
1585
+ """
1586
+ Specifies the times when the flow should be run when running on a
1587
+ production scheduler.
1588
+
1589
+ Parameters
1590
+ ----------
1591
+ hourly : bool, default False
1592
+ Run the workflow hourly.
1593
+ daily : bool, default True
1594
+ Run the workflow daily.
1595
+ weekly : bool, default False
1596
+ Run the workflow weekly.
1597
+ cron : str, optional, default None
1598
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1599
+ specified by this expression.
1600
+ timezone : str, optional, default None
1601
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1602
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1603
+ """
1604
+ ...
1605
+
1606
+ 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]]:
1607
+ """
1608
+ 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.
1609
+ 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.
1610
+
1611
+ Parameters
1612
+ ----------
1613
+ timeout : int
1614
+ Time, in seconds before the task times out and fails. (Default: 3600)
1615
+ poke_interval : int
1616
+ Time in seconds that the job should wait in between each try. (Default: 60)
1617
+ mode : str
1618
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1619
+ exponential_backoff : bool
1620
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1621
+ pool : str
1622
+ the slot pool this task should run in,
1623
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1624
+ soft_fail : bool
1625
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1626
+ name : str
1627
+ Name of the sensor on Airflow
1628
+ description : str
1629
+ Description of sensor in the Airflow UI
1630
+ external_dag_id : str
1631
+ The dag_id that contains the task you want to wait for.
1632
+ external_task_ids : List[str]
1633
+ The list of task_ids that you want to wait for.
1634
+ If None (default value) the sensor waits for the DAG. (Default: None)
1635
+ allowed_states : List[str]
1636
+ Iterable of allowed states, (Default: ['success'])
1637
+ failed_states : List[str]
1638
+ Iterable of failed or dis-allowed states. (Default: None)
1639
+ execution_delta : datetime.timedelta
1640
+ time difference with the previous execution to look at,
1641
+ the default is the same logical date as the current task or DAG. (Default: None)
1642
+ check_existence: bool
1643
+ Set to True to check if the external task exists or check if
1644
+ the DAG to wait for exists. (Default: True)
1645
+ """
1646
+ ...
1647
+
1648
+ def project(*, name: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1649
+ """
1650
+ Specifies what flows belong to the same project.
1651
+
1652
+ A project-specific namespace is created for all flows that
1653
+ use the same `@project(name)`.
1654
+
1655
+ Parameters
1656
+ ----------
1657
+ name : str
1658
+ Project name. Make sure that the name is unique amongst all
1659
+ projects that use the same production scheduler. The name may
1660
+ contain only lowercase alphanumeric characters and underscores.
1661
+
1662
+
1663
+ """
1664
+ ...
1665
+
1666
+ @typing.overload
1667
+ 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]]:
1668
+ """
1669
+ Specifies the Conda environment for all steps of the flow.
1670
+
1671
+ Use `@conda_base` to set common libraries required by all
1672
+ steps and use `@conda` to specify step-specific additions.
1673
+
1674
+ Parameters
1675
+ ----------
1676
+ packages : Dict[str, str], default {}
1707
1677
  Packages to use for this flow. The key is the name of the package
1708
1678
  and the value is the version to use.
1709
- python : str, optional, default: None
1679
+ libraries : Dict[str, str], default {}
1680
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1681
+ python : str, optional, default None
1710
1682
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
1711
1683
  that the version used will correspond to the version of the Python interpreter used to start the run.
1684
+ disabled : bool, default False
1685
+ If set to True, disables Conda.
1712
1686
  """
1713
1687
  ...
1714
1688
 
1715
1689
  @typing.overload
1716
- def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1690
+ def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1717
1691
  ...
1718
1692
 
1719
- def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1693
+ 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):
1720
1694
  """
1721
- Specifies the PyPI packages for all steps of the flow.
1695
+ Specifies the Conda environment for all steps of the flow.
1696
+
1697
+ Use `@conda_base` to set common libraries required by all
1698
+ steps and use `@conda` to specify step-specific additions.
1722
1699
 
1723
- Use `@pypi_base` to set common packages required by all
1724
- steps and use `@pypi` to specify step-specific overrides.
1725
1700
  Parameters
1726
1701
  ----------
1727
- packages : Dict[str, str], default: {}
1702
+ packages : Dict[str, str], default {}
1728
1703
  Packages to use for this flow. The key is the name of the package
1729
1704
  and the value is the version to use.
1730
- python : str, optional, default: None
1705
+ libraries : Dict[str, str], default {}
1706
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1707
+ python : str, optional, default None
1731
1708
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
1732
1709
  that the version used will correspond to the version of the Python interpreter used to start the run.
1710
+ disabled : bool, default False
1711
+ If set to True, disables Conda.
1733
1712
  """
1734
1713
  ...
1735
1714
 
@@ -1828,21 +1807,45 @@ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: t
1828
1807
  """
1829
1808
  ...
1830
1809
 
1831
- def project(*, name: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1810
+ 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]]:
1832
1811
  """
1833
- Specifies what flows belong to the same project.
1834
-
1835
- A project-specific namespace is created for all flows that
1836
- use the same `@project(name)`.
1812
+ 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)
1813
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1814
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1815
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
1816
+ starts only after all sensors finish.
1837
1817
 
1838
1818
  Parameters
1839
1819
  ----------
1820
+ timeout : int
1821
+ Time, in seconds before the task times out and fails. (Default: 3600)
1822
+ poke_interval : int
1823
+ Time in seconds that the job should wait in between each try. (Default: 60)
1824
+ mode : str
1825
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1826
+ exponential_backoff : bool
1827
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1828
+ pool : str
1829
+ the slot pool this task should run in,
1830
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1831
+ soft_fail : bool
1832
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1840
1833
  name : str
1841
- Project name. Make sure that the name is unique amongst all
1842
- projects that use the same production scheduler. The name may
1843
- contain only lowercase alphanumeric characters and underscores.
1844
-
1845
-
1834
+ Name of the sensor on Airflow
1835
+ description : str
1836
+ Description of sensor in the Airflow UI
1837
+ bucket_key : Union[str, List[str]]
1838
+ The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1839
+ When it's specified as a full s3:// url, please leave `bucket_name` as None
1840
+ bucket_name : str
1841
+ Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1842
+ When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1843
+ wildcard_match : bool
1844
+ whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1845
+ aws_conn_id : str
1846
+ a reference to the s3 connection on Airflow. (Default: None)
1847
+ verify : bool
1848
+ Whether or not to verify SSL certificates for S3 connection. (Default: None)
1846
1849
  """
1847
1850
  ...
1848
1851