metaflow-stubs 2.12.9__py2.py3-none-any.whl → 2.12.10__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 +602 -602
  2. metaflow-stubs/cards.pyi +6 -6
  3. metaflow-stubs/cli.pyi +2 -2
  4. metaflow-stubs/client/__init__.pyi +3 -3
  5. metaflow-stubs/client/core.pyi +7 -7
  6. metaflow-stubs/client/filecache.pyi +3 -3
  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 +5 -5
  13. metaflow-stubs/metadata/metadata.pyi +2 -2
  14. metaflow-stubs/metadata/util.pyi +2 -2
  15. metaflow-stubs/metaflow_config.pyi +2 -2
  16. metaflow-stubs/metaflow_current.pyi +18 -18
  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 +3 -3
  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 +16 -5
  35. metaflow-stubs/plugins/argo/argo_workflows_cli.pyi +8 -8
  36. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +7 -7
  37. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +4 -4
  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 +2 -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 +3 -3
  46. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +2 -2
  47. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +3 -3
  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 +2 -2
  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 +3 -3
  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 +4 -4
  67. metaflow-stubs/plugins/cards/card_client.pyi +3 -3
  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 +3 -3
  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 +7 -7
  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 +3 -3
  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 +4 -4
  108. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +2 -2
  109. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +3 -3
  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 +4 -4
  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 +2 -2
  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 +2 -2
  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 +4 -4
  129. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +3 -3
  130. metaflow-stubs/plugins/timeout_decorator.pyi +3 -3
  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 +5 -5
  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 +3 -3
  142. metaflow-stubs/system/system_monitor.pyi +2 -2
  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.9.dist-info → metaflow_stubs-2.12.10.dist-info}/METADATA +2 -2
  147. metaflow_stubs-2.12.10.dist-info/RECORD +150 -0
  148. {metaflow_stubs-2.12.9.dist-info → metaflow_stubs-2.12.10.dist-info}/WHEEL +1 -1
  149. metaflow_stubs-2.12.9.dist-info/RECORD +0 -150
  150. {metaflow_stubs-2.12.9.dist-info → metaflow_stubs-2.12.10.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.9 #
4
- # Generated on 2024-07-25T18:20:16.223359 #
3
+ # MF version: 2.12.10 #
4
+ # Generated on 2024-07-30T23:31:35.098122 #
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.runner.metaflow_runner
12
- import io
13
- import metaflow.flowspec
14
- import typing
15
- import datetime
16
- import metaflow.client.core
17
- import metaflow.parameters
18
11
  import metaflow.metaflow_current
19
- import metaflow.datastore.inputs
20
12
  import metaflow._vendor.click.types
13
+ import metaflow.client.core
21
14
  import metaflow.events
15
+ import metaflow.parameters
16
+ import datetime
17
+ import typing
22
18
  import metaflow.plugins.datatools.s3.s3
19
+ import metaflow.datastore.inputs
20
+ import io
21
+ import metaflow.flowspec
22
+ import metaflow.runner.metaflow_runner
23
23
  FlowSpecDerived = typing.TypeVar("FlowSpecDerived", bound="FlowSpec", contravariant=False, covariant=False)
24
24
  StepFlag = typing.NewType("StepFlag", bool)
25
25
 
@@ -727,65 +727,6 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
727
727
  """
728
728
  ...
729
729
 
730
- 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]]]:
731
- """
732
- Specifies that this step should execute on Kubernetes.
733
-
734
- Parameters
735
- ----------
736
- cpu : int, default 1
737
- Number of CPUs required for this step. If `@resources` is
738
- also present, the maximum value from all decorators is used.
739
- memory : int, default 4096
740
- Memory size (in MB) required for this step. If
741
- `@resources` is also present, the maximum value from all decorators is
742
- used.
743
- disk : int, default 10240
744
- Disk size (in MB) required for this step. If
745
- `@resources` is also present, the maximum value from all decorators is
746
- used.
747
- image : str, optional, default None
748
- Docker image to use when launching on Kubernetes. If not specified, and
749
- METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
750
- not, a default Docker image mapping to the current version of Python is used.
751
- image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
752
- If given, the imagePullPolicy to be applied to the Docker image of the step.
753
- service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
754
- Kubernetes service account to use when launching pod in Kubernetes.
755
- secrets : List[str], optional, default None
756
- Kubernetes secrets to use when launching pod in Kubernetes. These
757
- secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
758
- in Metaflow configuration.
759
- namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
760
- Kubernetes namespace to use when launching pod in Kubernetes.
761
- gpu : int, optional, default None
762
- Number of GPUs required for this step. A value of zero implies that
763
- the scheduled node should not have GPUs.
764
- gpu_vendor : str, default KUBERNETES_GPU_VENDOR
765
- The vendor of the GPUs to be used for this step.
766
- tolerations : List[str], default []
767
- The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
768
- Kubernetes tolerations to use when launching pod in Kubernetes.
769
- use_tmpfs : bool, default False
770
- This enables an explicit tmpfs mount for this step.
771
- tmpfs_tempdir : bool, default True
772
- sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
773
- tmpfs_size : int, optional, default: None
774
- The value for the size (in MiB) of the tmpfs mount for this step.
775
- This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
776
- memory allocated for this step.
777
- tmpfs_path : str, optional, default /metaflow_temp
778
- Path to tmpfs mount for this step.
779
- persistent_volume_claims : Dict[str, str], optional, default None
780
- A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
781
- volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
782
- shared_memory: int, optional
783
- Shared memory size (in MiB) required for this step
784
- port: int, optional
785
- Port number to specify in the Kubernetes job object
786
- """
787
- ...
788
-
789
730
  @typing.overload
790
731
  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]]]:
791
732
  """
@@ -838,333 +779,286 @@ def card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typ
838
779
  ...
839
780
 
840
781
  @typing.overload
841
- 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]]]:
782
+ 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]]]:
842
783
  """
843
- Specifies that this step should execute on [AWS Batch](https://aws.amazon.com/batch/).
784
+ Specifies environment variables to be set prior to the execution of a step.
844
785
 
845
786
  Parameters
846
787
  ----------
847
- cpu : int, default 1
848
- Number of CPUs required for this step. If `@resources` is
849
- also present, the maximum value from all decorators is used.
850
- gpu : int, default 0
851
- Number of GPUs required for this step. If `@resources` is
852
- also present, the maximum value from all decorators is used.
853
- memory : int, default 4096
854
- Memory size (in MB) required for this step. If
855
- `@resources` is also present, the maximum value from all decorators is
856
- used.
857
- image : str, optional, default None
858
- Docker image to use when launching on AWS Batch. If not specified, and
859
- METAFLOW_BATCH_CONTAINER_IMAGE is specified, that image is used. If
860
- not, a default Docker image mapping to the current version of Python is used.
861
- queue : str, default METAFLOW_BATCH_JOB_QUEUE
862
- AWS Batch Job Queue to submit the job to.
863
- iam_role : str, default METAFLOW_ECS_S3_ACCESS_IAM_ROLE
864
- AWS IAM role that AWS Batch container uses to access AWS cloud resources.
865
- execution_role : str, default METAFLOW_ECS_FARGATE_EXECUTION_ROLE
866
- AWS IAM role that AWS Batch can use [to trigger AWS Fargate tasks]
867
- (https://docs.aws.amazon.com/batch/latest/userguide/execution-IAM-role.html).
868
- shared_memory : int, optional, default None
869
- The value for the size (in MiB) of the /dev/shm volume for this step.
870
- This parameter maps to the `--shm-size` option in Docker.
871
- max_swap : int, optional, default None
872
- The total amount of swap memory (in MiB) a container can use for this
873
- step. This parameter is translated to the `--memory-swap` option in
874
- Docker where the value is the sum of the container memory plus the
875
- `max_swap` value.
876
- swappiness : int, optional, default None
877
- This allows you to tune memory swappiness behavior for this step.
878
- A swappiness value of 0 causes swapping not to happen unless absolutely
879
- necessary. A swappiness value of 100 causes pages to be swapped very
880
- aggressively. Accepted values are whole numbers between 0 and 100.
881
- use_tmpfs : bool, default False
882
- This enables an explicit tmpfs mount for this step. Note that tmpfs is
883
- not available on Fargate compute environments
884
- tmpfs_tempdir : bool, default True
885
- sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
886
- tmpfs_size : int, optional, default None
887
- The value for the size (in MiB) of the tmpfs mount for this step.
888
- This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
889
- memory allocated for this step.
890
- tmpfs_path : str, optional, default None
891
- Path to tmpfs mount for this step. Defaults to /metaflow_temp.
892
- inferentia : int, default 0
893
- Number of Inferentia chips required for this step.
894
- trainium : int, default None
895
- Alias for inferentia. Use only one of the two.
896
- efa : int, default 0
897
- Number of elastic fabric adapter network devices to attach to container
898
- ephemeral_storage : int, default None
899
- The total amount, in GiB, of ephemeral storage to set for the task, 21-200GiB.
900
- This is only relevant for Fargate compute environments
901
- log_driver: str, optional, default None
902
- The log driver to use for the Amazon ECS container.
903
- log_options: List[str], optional, default None
904
- List of strings containing options for the chosen log driver. The configurable values
905
- depend on the `log driver` chosen. Validation of these options is not supported yet.
906
- Example: [`awslogs-group:aws/batch/job`]
788
+ vars : Dict[str, str], default {}
789
+ Dictionary of environment variables to set.
907
790
  """
908
791
  ...
909
792
 
910
793
  @typing.overload
911
- def batch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
794
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
912
795
  ...
913
796
 
914
797
  @typing.overload
915
- def batch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
798
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
916
799
  ...
917
800
 
918
- 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):
801
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
919
802
  """
920
- Specifies that this step should execute on [AWS Batch](https://aws.amazon.com/batch/).
803
+ Specifies environment variables to be set prior to the execution of a step.
921
804
 
922
805
  Parameters
923
806
  ----------
924
- cpu : int, default 1
925
- Number of CPUs required for this step. If `@resources` is
926
- also present, the maximum value from all decorators is used.
927
- gpu : int, default 0
928
- Number of GPUs required for this step. If `@resources` is
929
- also present, the maximum value from all decorators is used.
930
- memory : int, default 4096
931
- Memory size (in MB) required for this step. If
932
- `@resources` is also present, the maximum value from all decorators is
933
- used.
934
- image : str, optional, default None
935
- Docker image to use when launching on AWS Batch. If not specified, and
936
- METAFLOW_BATCH_CONTAINER_IMAGE is specified, that image is used. If
937
- not, a default Docker image mapping to the current version of Python is used.
938
- queue : str, default METAFLOW_BATCH_JOB_QUEUE
939
- AWS Batch Job Queue to submit the job to.
940
- iam_role : str, default METAFLOW_ECS_S3_ACCESS_IAM_ROLE
941
- AWS IAM role that AWS Batch container uses to access AWS cloud resources.
942
- execution_role : str, default METAFLOW_ECS_FARGATE_EXECUTION_ROLE
943
- AWS IAM role that AWS Batch can use [to trigger AWS Fargate tasks]
944
- (https://docs.aws.amazon.com/batch/latest/userguide/execution-IAM-role.html).
945
- shared_memory : int, optional, default None
946
- The value for the size (in MiB) of the /dev/shm volume for this step.
947
- This parameter maps to the `--shm-size` option in Docker.
948
- max_swap : int, optional, default None
949
- The total amount of swap memory (in MiB) a container can use for this
950
- step. This parameter is translated to the `--memory-swap` option in
951
- Docker where the value is the sum of the container memory plus the
952
- `max_swap` value.
953
- swappiness : int, optional, default None
954
- This allows you to tune memory swappiness behavior for this step.
955
- A swappiness value of 0 causes swapping not to happen unless absolutely
956
- necessary. A swappiness value of 100 causes pages to be swapped very
957
- aggressively. Accepted values are whole numbers between 0 and 100.
958
- use_tmpfs : bool, default False
959
- This enables an explicit tmpfs mount for this step. Note that tmpfs is
960
- not available on Fargate compute environments
961
- tmpfs_tempdir : bool, default True
962
- sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
963
- tmpfs_size : int, optional, default None
964
- The value for the size (in MiB) of the tmpfs mount for this step.
965
- This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
966
- memory allocated for this step.
967
- tmpfs_path : str, optional, default None
968
- Path to tmpfs mount for this step. Defaults to /metaflow_temp.
969
- inferentia : int, default 0
970
- Number of Inferentia chips required for this step.
971
- trainium : int, default None
972
- Alias for inferentia. Use only one of the two.
973
- efa : int, default 0
974
- Number of elastic fabric adapter network devices to attach to container
975
- ephemeral_storage : int, default None
976
- The total amount, in GiB, of ephemeral storage to set for the task, 21-200GiB.
977
- This is only relevant for Fargate compute environments
978
- log_driver: str, optional, default None
979
- The log driver to use for the Amazon ECS container.
980
- log_options: List[str], optional, default None
981
- List of strings containing options for the chosen log driver. The configurable values
982
- depend on the `log driver` chosen. Validation of these options is not supported yet.
983
- Example: [`awslogs-group:aws/batch/job`]
807
+ vars : Dict[str, str], default {}
808
+ Dictionary of environment variables to set.
984
809
  """
985
810
  ...
986
811
 
987
812
  @typing.overload
988
- 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]]]:
813
+ 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]]]:
989
814
  """
990
- Specifies the PyPI packages for the step.
815
+ Specifies that the step will success under all circumstances.
991
816
 
992
- Information in this decorator will augment any
993
- attributes set in the `@pyi_base` flow-level decorator. Hence,
994
- you can use `@pypi_base` to set packages required by all
995
- steps and use `@pypi` to specify step-specific overrides.
817
+ The decorator will create an optional artifact, specified by `var`, which
818
+ contains the exception raised. You can use it to detect the presence
819
+ of errors, indicating that all happy-path artifacts produced by the step
820
+ are missing.
996
821
 
997
822
  Parameters
998
823
  ----------
999
- packages : Dict[str, str], default: {}
1000
- Packages to use for this step. The key is the name of the package
1001
- and the value is the version to use.
1002
- python : str, optional, default: None
1003
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1004
- that the version used will correspond to the version of the Python interpreter used to start the run.
824
+ var : str, optional, default None
825
+ Name of the artifact in which to store the caught exception.
826
+ If not specified, the exception is not stored.
827
+ print_exception : bool, default True
828
+ Determines whether or not the exception is printed to
829
+ stdout when caught.
1005
830
  """
1006
831
  ...
1007
832
 
1008
833
  @typing.overload
1009
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
834
+ def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1010
835
  ...
1011
836
 
1012
837
  @typing.overload
1013
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
838
+ def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1014
839
  ...
1015
840
 
1016
- 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):
841
+ 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):
1017
842
  """
1018
- Specifies the PyPI packages for the step.
843
+ Specifies that the step will success under all circumstances.
1019
844
 
1020
- Information in this decorator will augment any
1021
- attributes set in the `@pyi_base` flow-level decorator. Hence,
1022
- you can use `@pypi_base` to set packages required by all
1023
- steps and use `@pypi` to specify step-specific overrides.
845
+ The decorator will create an optional artifact, specified by `var`, which
846
+ contains the exception raised. You can use it to detect the presence
847
+ of errors, indicating that all happy-path artifacts produced by the step
848
+ are missing.
1024
849
 
1025
850
  Parameters
1026
851
  ----------
1027
- packages : Dict[str, str], default: {}
1028
- Packages to use for this step. The key is the name of the package
1029
- and the value is the version to use.
1030
- python : str, optional, default: None
1031
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1032
- that the version used will correspond to the version of the Python interpreter used to start the run.
852
+ var : str, optional, default None
853
+ Name of the artifact in which to store the caught exception.
854
+ If not specified, the exception is not stored.
855
+ print_exception : bool, default True
856
+ Determines whether or not the exception is printed to
857
+ stdout when caught.
1033
858
  """
1034
859
  ...
1035
860
 
1036
861
  @typing.overload
1037
- 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]]]:
862
+ 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]]]:
1038
863
  """
1039
- Specifies that the step will success under all circumstances.
864
+ Specifies the PyPI packages for the step.
1040
865
 
1041
- The decorator will create an optional artifact, specified by `var`, which
1042
- contains the exception raised. You can use it to detect the presence
1043
- of errors, indicating that all happy-path artifacts produced by the step
1044
- are missing.
866
+ Information in this decorator will augment any
867
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
868
+ you can use `@pypi_base` to set packages required by all
869
+ steps and use `@pypi` to specify step-specific overrides.
1045
870
 
1046
871
  Parameters
1047
872
  ----------
1048
- var : str, optional, default None
1049
- Name of the artifact in which to store the caught exception.
1050
- If not specified, the exception is not stored.
1051
- print_exception : bool, default True
1052
- Determines whether or not the exception is printed to
1053
- stdout when caught.
873
+ packages : Dict[str, str], default: {}
874
+ Packages to use for this step. The key is the name of the package
875
+ and the value is the version to use.
876
+ python : str, optional, default: None
877
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
878
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1054
879
  """
1055
880
  ...
1056
881
 
1057
882
  @typing.overload
1058
- def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
883
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1059
884
  ...
1060
885
 
1061
886
  @typing.overload
1062
- def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
887
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1063
888
  ...
1064
889
 
1065
- 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):
890
+ 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):
1066
891
  """
1067
- Specifies that the step will success under all circumstances.
892
+ Specifies the PyPI packages for the step.
1068
893
 
1069
- The decorator will create an optional artifact, specified by `var`, which
1070
- contains the exception raised. You can use it to detect the presence
1071
- of errors, indicating that all happy-path artifacts produced by the step
1072
- are missing.
894
+ Information in this decorator will augment any
895
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
896
+ you can use `@pypi_base` to set packages required by all
897
+ steps and use `@pypi` to specify step-specific overrides.
1073
898
 
1074
899
  Parameters
1075
900
  ----------
1076
- var : str, optional, default None
1077
- Name of the artifact in which to store the caught exception.
1078
- If not specified, the exception is not stored.
1079
- print_exception : bool, default True
1080
- Determines whether or not the exception is printed to
1081
- stdout when caught.
901
+ packages : Dict[str, str], default: {}
902
+ Packages to use for this step. The key is the name of the package
903
+ and the value is the version to use.
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.
1082
907
  """
1083
908
  ...
1084
909
 
1085
910
  @typing.overload
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]]]:
911
+ 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]]]:
1087
912
  """
1088
- Specifies secrets to be retrieved and injected as environment variables prior to
1089
- the execution of a step.
913
+ Specifies the resources needed when executing this step.
914
+
915
+ Use `@resources` to specify the resource requirements
916
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
917
+
918
+ You can choose the compute layer on the command line by executing e.g.
919
+ ```
920
+ python myflow.py run --with batch
921
+ ```
922
+ or
923
+ ```
924
+ python myflow.py run --with kubernetes
925
+ ```
926
+ which executes the flow on the desired system using the
927
+ requirements specified in `@resources`.
1090
928
 
1091
929
  Parameters
1092
930
  ----------
1093
- sources : List[Union[str, Dict[str, Any]]], default: []
1094
- List of secret specs, defining how the secrets are to be retrieved
931
+ cpu : int, default 1
932
+ Number of CPUs required for this step.
933
+ gpu : int, default 0
934
+ Number of GPUs required for this step.
935
+ disk : int, optional, default None
936
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
937
+ memory : int, default 4096
938
+ Memory size (in MB) required for this step.
939
+ shared_memory : int, optional, default None
940
+ The value for the size (in MiB) of the /dev/shm volume for this step.
941
+ This parameter maps to the `--shm-size` option in Docker.
1095
942
  """
1096
943
  ...
1097
944
 
1098
945
  @typing.overload
1099
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
946
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1100
947
  ...
1101
948
 
1102
949
  @typing.overload
1103
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
950
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1104
951
  ...
1105
952
 
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]]] = []):
953
+ 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):
1107
954
  """
1108
- Specifies secrets to be retrieved and injected as environment variables prior to
1109
- the execution of a step.
955
+ Specifies the resources needed when executing this step.
956
+
957
+ Use `@resources` to specify the resource requirements
958
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
959
+
960
+ You can choose the compute layer on the command line by executing e.g.
961
+ ```
962
+ python myflow.py run --with batch
963
+ ```
964
+ or
965
+ ```
966
+ python myflow.py run --with kubernetes
967
+ ```
968
+ which executes the flow on the desired system using the
969
+ requirements specified in `@resources`.
1110
970
 
1111
971
  Parameters
1112
972
  ----------
1113
- sources : List[Union[str, Dict[str, Any]]], default: []
1114
- List of secret specs, defining how the secrets are to be retrieved
973
+ cpu : int, default 1
974
+ Number of CPUs required for this step.
975
+ gpu : int, default 0
976
+ Number of GPUs required for this step.
977
+ disk : int, optional, default None
978
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
979
+ memory : int, default 4096
980
+ Memory size (in MB) required for this step.
981
+ shared_memory : int, optional, default None
982
+ The value for the size (in MiB) of the /dev/shm volume for this step.
983
+ This parameter maps to the `--shm-size` option in Docker.
1115
984
  """
1116
985
  ...
1117
986
 
1118
987
  @typing.overload
1119
- 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]]]:
988
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1120
989
  """
1121
- Specifies the number of times the task corresponding
1122
- to a step needs to be retried.
1123
-
1124
- This decorator is useful for handling transient errors, such as networking issues.
1125
- If your task contains operations that can't be retried safely, e.g. database updates,
1126
- it is advisable to annotate it with `@retry(times=0)`.
1127
-
1128
- This can be used in conjunction with the `@catch` decorator. The `@catch`
1129
- decorator will execute a no-op task after all retries have been exhausted,
1130
- ensuring that the flow execution can continue.
1131
-
1132
- Parameters
1133
- ----------
1134
- times : int, default 3
1135
- Number of times to retry this task.
1136
- minutes_between_retries : int, default 2
1137
- Number of minutes between retries.
990
+ Decorator prototype for all step decorators. This function gets specialized
991
+ and imported for all decorators types by _import_plugin_decorators().
1138
992
  """
1139
993
  ...
1140
994
 
1141
995
  @typing.overload
1142
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
996
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1143
997
  ...
1144
998
 
1145
- @typing.overload
1146
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
999
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1000
+ """
1001
+ Decorator prototype for all step decorators. This function gets specialized
1002
+ and imported for all decorators types by _import_plugin_decorators().
1003
+ """
1147
1004
  ...
1148
1005
 
1149
- 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):
1006
+ 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]]]:
1150
1007
  """
1151
- Specifies the number of times the task corresponding
1152
- to a step needs to be retried.
1153
-
1154
- This decorator is useful for handling transient errors, such as networking issues.
1155
- If your task contains operations that can't be retried safely, e.g. database updates,
1156
- it is advisable to annotate it with `@retry(times=0)`.
1157
-
1158
- This can be used in conjunction with the `@catch` decorator. The `@catch`
1159
- decorator will execute a no-op task after all retries have been exhausted,
1160
- ensuring that the flow execution can continue.
1008
+ Specifies that this step should execute on Kubernetes.
1161
1009
 
1162
1010
  Parameters
1163
1011
  ----------
1164
- times : int, default 3
1165
- Number of times to retry this task.
1166
- minutes_between_retries : int, default 2
1167
- Number of minutes between retries.
1012
+ cpu : int, default 1
1013
+ Number of CPUs required for this step. If `@resources` is
1014
+ also present, the maximum value from all decorators is used.
1015
+ memory : int, default 4096
1016
+ Memory size (in MB) required for this step. If
1017
+ `@resources` is also present, the maximum value from all decorators is
1018
+ used.
1019
+ disk : int, default 10240
1020
+ Disk size (in MB) required for this step. If
1021
+ `@resources` is also present, the maximum value from all decorators is
1022
+ used.
1023
+ image : str, optional, default None
1024
+ Docker image to use when launching on Kubernetes. If not specified, and
1025
+ METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
1026
+ not, a default Docker image mapping to the current version of Python is used.
1027
+ image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
1028
+ If given, the imagePullPolicy to be applied to the Docker image of the step.
1029
+ service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
1030
+ Kubernetes service account to use when launching pod in Kubernetes.
1031
+ secrets : List[str], optional, default None
1032
+ Kubernetes secrets to use when launching pod in Kubernetes. These
1033
+ secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
1034
+ in Metaflow configuration.
1035
+ namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
1036
+ Kubernetes namespace to use when launching pod in Kubernetes.
1037
+ gpu : int, optional, default None
1038
+ Number of GPUs required for this step. A value of zero implies that
1039
+ the scheduled node should not have GPUs.
1040
+ gpu_vendor : str, default KUBERNETES_GPU_VENDOR
1041
+ The vendor of the GPUs to be used for this step.
1042
+ tolerations : List[str], default []
1043
+ The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
1044
+ Kubernetes tolerations to use when launching pod in Kubernetes.
1045
+ use_tmpfs : bool, default False
1046
+ This enables an explicit tmpfs mount for this step.
1047
+ tmpfs_tempdir : bool, default True
1048
+ sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
1049
+ tmpfs_size : int, optional, default: None
1050
+ The value for the size (in MiB) of the tmpfs mount for this step.
1051
+ This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
1052
+ memory allocated for this step.
1053
+ tmpfs_path : str, optional, default /metaflow_temp
1054
+ Path to tmpfs mount for this step.
1055
+ persistent_volume_claims : Dict[str, str], optional, default None
1056
+ A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
1057
+ volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
1058
+ shared_memory: int, optional
1059
+ Shared memory size (in MiB) required for this step
1060
+ port: int, optional
1061
+ Port number to specify in the Kubernetes job object
1168
1062
  """
1169
1063
  ...
1170
1064
 
@@ -1226,240 +1120,364 @@ def conda(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
1226
1120
  ...
1227
1121
 
1228
1122
  @typing.overload
1229
- 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]]]:
1123
+ 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]]]:
1230
1124
  """
1231
- Specifies environment variables to be set prior to the execution of a step.
1125
+ Specifies a timeout for your step.
1126
+
1127
+ This decorator is useful if this step may hang indefinitely.
1128
+
1129
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
1130
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
1131
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
1132
+
1133
+ Note that all the values specified in parameters are added together so if you specify
1134
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
1232
1135
 
1233
1136
  Parameters
1234
1137
  ----------
1235
- vars : Dict[str, str], default {}
1236
- Dictionary of environment variables to set.
1138
+ seconds : int, default 0
1139
+ Number of seconds to wait prior to timing out.
1140
+ minutes : int, default 0
1141
+ Number of minutes to wait prior to timing out.
1142
+ hours : int, default 0
1143
+ Number of hours to wait prior to timing out.
1237
1144
  """
1238
1145
  ...
1239
1146
 
1240
1147
  @typing.overload
1241
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1148
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1242
1149
  ...
1243
1150
 
1244
1151
  @typing.overload
1245
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1152
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1246
1153
  ...
1247
1154
 
1248
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
1155
+ 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):
1249
1156
  """
1250
- Specifies environment variables to be set prior to the execution of a step.
1157
+ Specifies a timeout for your step.
1251
1158
 
1252
- Parameters
1253
- ----------
1254
- vars : Dict[str, str], default {}
1255
- Dictionary of environment variables to set.
1159
+ This decorator is useful if this step may hang indefinitely.
1160
+
1161
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
1162
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
1163
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
1164
+
1165
+ Note that all the values specified in parameters are added together so if you specify
1166
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
1167
+
1168
+ Parameters
1169
+ ----------
1170
+ seconds : int, default 0
1171
+ Number of seconds to wait prior to timing out.
1172
+ minutes : int, default 0
1173
+ Number of minutes to wait prior to timing out.
1174
+ hours : int, default 0
1175
+ Number of hours to wait prior to timing out.
1256
1176
  """
1257
1177
  ...
1258
1178
 
1259
1179
  @typing.overload
1260
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1180
+ 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]]]:
1261
1181
  """
1262
- Decorator prototype for all step decorators. This function gets specialized
1263
- and imported for all decorators types by _import_plugin_decorators().
1182
+ Specifies that this step should execute on [AWS Batch](https://aws.amazon.com/batch/).
1183
+
1184
+ Parameters
1185
+ ----------
1186
+ cpu : int, default 1
1187
+ Number of CPUs required for this step. If `@resources` is
1188
+ also present, the maximum value from all decorators is used.
1189
+ gpu : int, default 0
1190
+ Number of GPUs required for this step. If `@resources` is
1191
+ also present, the maximum value from all decorators is used.
1192
+ memory : int, default 4096
1193
+ Memory size (in MB) required for this step. If
1194
+ `@resources` is also present, the maximum value from all decorators is
1195
+ used.
1196
+ image : str, optional, default None
1197
+ Docker image to use when launching on AWS Batch. If not specified, and
1198
+ METAFLOW_BATCH_CONTAINER_IMAGE is specified, that image is used. If
1199
+ not, a default Docker image mapping to the current version of Python is used.
1200
+ queue : str, default METAFLOW_BATCH_JOB_QUEUE
1201
+ AWS Batch Job Queue to submit the job to.
1202
+ iam_role : str, default METAFLOW_ECS_S3_ACCESS_IAM_ROLE
1203
+ AWS IAM role that AWS Batch container uses to access AWS cloud resources.
1204
+ execution_role : str, default METAFLOW_ECS_FARGATE_EXECUTION_ROLE
1205
+ AWS IAM role that AWS Batch can use [to trigger AWS Fargate tasks]
1206
+ (https://docs.aws.amazon.com/batch/latest/userguide/execution-IAM-role.html).
1207
+ shared_memory : int, optional, default None
1208
+ The value for the size (in MiB) of the /dev/shm volume for this step.
1209
+ This parameter maps to the `--shm-size` option in Docker.
1210
+ max_swap : int, optional, default None
1211
+ The total amount of swap memory (in MiB) a container can use for this
1212
+ step. This parameter is translated to the `--memory-swap` option in
1213
+ Docker where the value is the sum of the container memory plus the
1214
+ `max_swap` value.
1215
+ swappiness : int, optional, default None
1216
+ This allows you to tune memory swappiness behavior for this step.
1217
+ A swappiness value of 0 causes swapping not to happen unless absolutely
1218
+ necessary. A swappiness value of 100 causes pages to be swapped very
1219
+ aggressively. Accepted values are whole numbers between 0 and 100.
1220
+ use_tmpfs : bool, default False
1221
+ This enables an explicit tmpfs mount for this step. Note that tmpfs is
1222
+ not available on Fargate compute environments
1223
+ tmpfs_tempdir : bool, default True
1224
+ sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
1225
+ tmpfs_size : int, optional, default None
1226
+ The value for the size (in MiB) of the tmpfs mount for this step.
1227
+ This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
1228
+ memory allocated for this step.
1229
+ tmpfs_path : str, optional, default None
1230
+ Path to tmpfs mount for this step. Defaults to /metaflow_temp.
1231
+ inferentia : int, default 0
1232
+ Number of Inferentia chips required for this step.
1233
+ trainium : int, default None
1234
+ Alias for inferentia. Use only one of the two.
1235
+ efa : int, default 0
1236
+ Number of elastic fabric adapter network devices to attach to container
1237
+ ephemeral_storage : int, default None
1238
+ The total amount, in GiB, of ephemeral storage to set for the task, 21-200GiB.
1239
+ This is only relevant for Fargate compute environments
1240
+ log_driver: str, optional, default None
1241
+ The log driver to use for the Amazon ECS container.
1242
+ log_options: List[str], optional, default None
1243
+ List of strings containing options for the chosen log driver. The configurable values
1244
+ depend on the `log driver` chosen. Validation of these options is not supported yet.
1245
+ Example: [`awslogs-group:aws/batch/job`]
1264
1246
  """
1265
1247
  ...
1266
1248
 
1267
1249
  @typing.overload
1268
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1250
+ def batch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1269
1251
  ...
1270
1252
 
1271
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1253
+ @typing.overload
1254
+ def batch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1255
+ ...
1256
+
1257
+ 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):
1272
1258
  """
1273
- Decorator prototype for all step decorators. This function gets specialized
1274
- and imported for all decorators types by _import_plugin_decorators().
1259
+ Specifies that this step should execute on [AWS Batch](https://aws.amazon.com/batch/).
1260
+
1261
+ Parameters
1262
+ ----------
1263
+ cpu : int, default 1
1264
+ Number of CPUs required for this step. If `@resources` is
1265
+ also present, the maximum value from all decorators is used.
1266
+ gpu : int, default 0
1267
+ Number of GPUs required for this step. If `@resources` is
1268
+ also present, the maximum value from all decorators is used.
1269
+ memory : int, default 4096
1270
+ Memory size (in MB) required for this step. If
1271
+ `@resources` is also present, the maximum value from all decorators is
1272
+ used.
1273
+ image : str, optional, default None
1274
+ Docker image to use when launching on AWS Batch. If not specified, and
1275
+ METAFLOW_BATCH_CONTAINER_IMAGE is specified, that image is used. If
1276
+ not, a default Docker image mapping to the current version of Python is used.
1277
+ queue : str, default METAFLOW_BATCH_JOB_QUEUE
1278
+ AWS Batch Job Queue to submit the job to.
1279
+ iam_role : str, default METAFLOW_ECS_S3_ACCESS_IAM_ROLE
1280
+ AWS IAM role that AWS Batch container uses to access AWS cloud resources.
1281
+ execution_role : str, default METAFLOW_ECS_FARGATE_EXECUTION_ROLE
1282
+ AWS IAM role that AWS Batch can use [to trigger AWS Fargate tasks]
1283
+ (https://docs.aws.amazon.com/batch/latest/userguide/execution-IAM-role.html).
1284
+ shared_memory : int, optional, default None
1285
+ The value for the size (in MiB) of the /dev/shm volume for this step.
1286
+ This parameter maps to the `--shm-size` option in Docker.
1287
+ max_swap : int, optional, default None
1288
+ The total amount of swap memory (in MiB) a container can use for this
1289
+ step. This parameter is translated to the `--memory-swap` option in
1290
+ Docker where the value is the sum of the container memory plus the
1291
+ `max_swap` value.
1292
+ swappiness : int, optional, default None
1293
+ This allows you to tune memory swappiness behavior for this step.
1294
+ A swappiness value of 0 causes swapping not to happen unless absolutely
1295
+ necessary. A swappiness value of 100 causes pages to be swapped very
1296
+ aggressively. Accepted values are whole numbers between 0 and 100.
1297
+ use_tmpfs : bool, default False
1298
+ This enables an explicit tmpfs mount for this step. Note that tmpfs is
1299
+ not available on Fargate compute environments
1300
+ tmpfs_tempdir : bool, default True
1301
+ sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
1302
+ tmpfs_size : int, optional, default None
1303
+ The value for the size (in MiB) of the tmpfs mount for this step.
1304
+ This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
1305
+ memory allocated for this step.
1306
+ tmpfs_path : str, optional, default None
1307
+ Path to tmpfs mount for this step. Defaults to /metaflow_temp.
1308
+ inferentia : int, default 0
1309
+ Number of Inferentia chips required for this step.
1310
+ trainium : int, default None
1311
+ Alias for inferentia. Use only one of the two.
1312
+ efa : int, default 0
1313
+ Number of elastic fabric adapter network devices to attach to container
1314
+ ephemeral_storage : int, default None
1315
+ The total amount, in GiB, of ephemeral storage to set for the task, 21-200GiB.
1316
+ This is only relevant for Fargate compute environments
1317
+ log_driver: str, optional, default None
1318
+ The log driver to use for the Amazon ECS container.
1319
+ log_options: List[str], optional, default None
1320
+ List of strings containing options for the chosen log driver. The configurable values
1321
+ depend on the `log driver` chosen. Validation of these options is not supported yet.
1322
+ Example: [`awslogs-group:aws/batch/job`]
1275
1323
  """
1276
1324
  ...
1277
1325
 
1278
1326
  @typing.overload
1279
- 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]]]:
1327
+ 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]]]:
1280
1328
  """
1281
- Specifies a timeout for your step.
1282
-
1283
- This decorator is useful if this step may hang indefinitely.
1284
-
1285
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
1286
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
1287
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
1288
-
1289
- Note that all the values specified in parameters are added together so if you specify
1290
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
1329
+ Specifies secrets to be retrieved and injected as environment variables prior to
1330
+ the execution of a step.
1291
1331
 
1292
1332
  Parameters
1293
1333
  ----------
1294
- seconds : int, default 0
1295
- Number of seconds to wait prior to timing out.
1296
- minutes : int, default 0
1297
- Number of minutes to wait prior to timing out.
1298
- hours : int, default 0
1299
- Number of hours to wait prior to timing out.
1334
+ sources : List[Union[str, Dict[str, Any]]], default: []
1335
+ List of secret specs, defining how the secrets are to be retrieved
1300
1336
  """
1301
1337
  ...
1302
1338
 
1303
1339
  @typing.overload
1304
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1340
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1305
1341
  ...
1306
1342
 
1307
1343
  @typing.overload
1308
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1344
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1309
1345
  ...
1310
1346
 
1311
- 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):
1347
+ 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]]] = []):
1312
1348
  """
1313
- Specifies a timeout for your step.
1314
-
1315
- This decorator is useful if this step may hang indefinitely.
1316
-
1317
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
1318
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
1319
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
1320
-
1321
- Note that all the values specified in parameters are added together so if you specify
1322
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
1349
+ Specifies secrets to be retrieved and injected as environment variables prior to
1350
+ the execution of a step.
1323
1351
 
1324
1352
  Parameters
1325
1353
  ----------
1326
- seconds : int, default 0
1327
- Number of seconds to wait prior to timing out.
1328
- minutes : int, default 0
1329
- Number of minutes to wait prior to timing out.
1330
- hours : int, default 0
1331
- Number of hours to wait prior to timing out.
1354
+ sources : List[Union[str, Dict[str, Any]]], default: []
1355
+ List of secret specs, defining how the secrets are to be retrieved
1332
1356
  """
1333
1357
  ...
1334
1358
 
1335
1359
  @typing.overload
1336
- 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]]]:
1360
+ 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]]]:
1337
1361
  """
1338
- Specifies the resources needed when executing this step.
1362
+ Specifies the number of times the task corresponding
1363
+ to a step needs to be retried.
1339
1364
 
1340
- Use `@resources` to specify the resource requirements
1341
- independently of the specific compute layer (`@batch`, `@kubernetes`).
1365
+ This decorator is useful for handling transient errors, such as networking issues.
1366
+ If your task contains operations that can't be retried safely, e.g. database updates,
1367
+ it is advisable to annotate it with `@retry(times=0)`.
1342
1368
 
1343
- You can choose the compute layer on the command line by executing e.g.
1344
- ```
1345
- python myflow.py run --with batch
1346
- ```
1347
- or
1348
- ```
1349
- python myflow.py run --with kubernetes
1350
- ```
1351
- which executes the flow on the desired system using the
1352
- requirements specified in `@resources`.
1369
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
1370
+ decorator will execute a no-op task after all retries have been exhausted,
1371
+ ensuring that the flow execution can continue.
1353
1372
 
1354
1373
  Parameters
1355
1374
  ----------
1356
- cpu : int, default 1
1357
- Number of CPUs required for this step.
1358
- gpu : int, default 0
1359
- Number of GPUs required for this step.
1360
- disk : int, optional, default None
1361
- Disk size (in MB) required for this step. Only applies on Kubernetes.
1362
- memory : int, default 4096
1363
- Memory size (in MB) required for this step.
1364
- shared_memory : int, optional, default None
1365
- The value for the size (in MiB) of the /dev/shm volume for this step.
1366
- This parameter maps to the `--shm-size` option in Docker.
1375
+ times : int, default 3
1376
+ Number of times to retry this task.
1377
+ minutes_between_retries : int, default 2
1378
+ Number of minutes between retries.
1367
1379
  """
1368
1380
  ...
1369
1381
 
1370
1382
  @typing.overload
1371
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1383
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1372
1384
  ...
1373
1385
 
1374
1386
  @typing.overload
1375
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1387
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1376
1388
  ...
1377
1389
 
1378
- 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):
1390
+ 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):
1379
1391
  """
1380
- Specifies the resources needed when executing this step.
1392
+ Specifies the number of times the task corresponding
1393
+ to a step needs to be retried.
1381
1394
 
1382
- Use `@resources` to specify the resource requirements
1383
- independently of the specific compute layer (`@batch`, `@kubernetes`).
1395
+ This decorator is useful for handling transient errors, such as networking issues.
1396
+ If your task contains operations that can't be retried safely, e.g. database updates,
1397
+ it is advisable to annotate it with `@retry(times=0)`.
1384
1398
 
1385
- You can choose the compute layer on the command line by executing e.g.
1386
- ```
1387
- python myflow.py run --with batch
1388
- ```
1389
- or
1390
- ```
1391
- python myflow.py run --with kubernetes
1392
- ```
1393
- which executes the flow on the desired system using the
1394
- requirements specified in `@resources`.
1399
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
1400
+ decorator will execute a no-op task after all retries have been exhausted,
1401
+ ensuring that the flow execution can continue.
1395
1402
 
1396
1403
  Parameters
1397
1404
  ----------
1398
- cpu : int, default 1
1399
- Number of CPUs required for this step.
1400
- gpu : int, default 0
1401
- Number of GPUs required for this step.
1402
- disk : int, optional, default None
1403
- Disk size (in MB) required for this step. Only applies on Kubernetes.
1404
- memory : int, default 4096
1405
- Memory size (in MB) required for this step.
1406
- shared_memory : int, optional, default None
1407
- The value for the size (in MiB) of the /dev/shm volume for this step.
1408
- This parameter maps to the `--shm-size` option in Docker.
1405
+ times : int, default 3
1406
+ Number of times to retry this task.
1407
+ minutes_between_retries : int, default 2
1408
+ Number of minutes between retries.
1409
1409
  """
1410
1410
  ...
1411
1411
 
1412
1412
  @typing.overload
1413
- 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]]:
1413
+ 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]]:
1414
1414
  """
1415
- Specifies the event(s) that this flow depends on.
1415
+ Specifies the times when the flow should be run when running on a
1416
+ production scheduler.
1416
1417
 
1417
- ```
1418
- @trigger(event='foo')
1419
- ```
1420
- or
1421
- ```
1422
- @trigger(events=['foo', 'bar'])
1423
- ```
1418
+ Parameters
1419
+ ----------
1420
+ hourly : bool, default False
1421
+ Run the workflow hourly.
1422
+ daily : bool, default True
1423
+ Run the workflow daily.
1424
+ weekly : bool, default False
1425
+ Run the workflow weekly.
1426
+ cron : str, optional, default None
1427
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1428
+ specified by this expression.
1429
+ timezone : str, optional, default None
1430
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1431
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1432
+ """
1433
+ ...
1434
+
1435
+ @typing.overload
1436
+ def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1437
+ ...
1438
+
1439
+ 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):
1440
+ """
1441
+ Specifies the times when the flow should be run when running on a
1442
+ production scheduler.
1424
1443
 
1425
- Additionally, you can specify the parameter mappings
1426
- to map event payload to Metaflow parameters for the flow.
1427
- ```
1428
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1429
- ```
1430
- or
1431
- ```
1432
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1433
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1434
- ```
1444
+ Parameters
1445
+ ----------
1446
+ hourly : bool, default False
1447
+ Run the workflow hourly.
1448
+ daily : bool, default True
1449
+ Run the workflow daily.
1450
+ weekly : bool, default False
1451
+ Run the workflow weekly.
1452
+ cron : str, optional, default None
1453
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1454
+ specified by this expression.
1455
+ timezone : str, optional, default None
1456
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1457
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1458
+ """
1459
+ ...
1460
+
1461
+ def project(*, name: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1462
+ """
1463
+ Specifies what flows belong to the same project.
1435
1464
 
1436
- 'parameters' can also be a list of strings and tuples like so:
1437
- ```
1438
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1439
- ```
1440
- This is equivalent to:
1441
- ```
1442
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1443
- ```
1465
+ A project-specific namespace is created for all flows that
1466
+ use the same `@project(name)`.
1444
1467
 
1445
1468
  Parameters
1446
1469
  ----------
1447
- event : Union[str, Dict[str, Any]], optional, default None
1448
- Event dependency for this flow.
1449
- events : List[Union[str, Dict[str, Any]]], default []
1450
- Events dependency for this flow.
1451
- options : Dict[str, Any], default {}
1452
- Backend-specific configuration for tuning eventing behavior.
1470
+ name : str
1471
+ Project name. Make sure that the name is unique amongst all
1472
+ projects that use the same production scheduler. The name may
1473
+ contain only lowercase alphanumeric characters and underscores.
1453
1474
 
1454
1475
 
1455
1476
  """
1456
1477
  ...
1457
1478
 
1458
1479
  @typing.overload
1459
- def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1460
- ...
1461
-
1462
- 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] = {}):
1480
+ 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]]:
1463
1481
  """
1464
1482
  Specifies the event(s) that this flow depends on.
1465
1483
 
@@ -1495,58 +1513,8 @@ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: t
1495
1513
  ----------
1496
1514
  event : Union[str, Dict[str, Any]], optional, default None
1497
1515
  Event dependency for this flow.
1498
- events : List[Union[str, Dict[str, Any]]], default []
1499
- Events dependency for this flow.
1500
- options : Dict[str, Any], default {}
1501
- Backend-specific configuration for tuning eventing behavior.
1502
-
1503
-
1504
- """
1505
- ...
1506
-
1507
- @typing.overload
1508
- def trigger_on_finish(*, flow: typing.Union[str, typing.Dict[str, str], None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1509
- """
1510
- Specifies the flow(s) that this flow depends on.
1511
-
1512
- ```
1513
- @trigger_on_finish(flow='FooFlow')
1514
- ```
1515
- or
1516
- ```
1517
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1518
- ```
1519
- This decorator respects the @project decorator and triggers the flow
1520
- when upstream runs within the same namespace complete successfully
1521
-
1522
- Additionally, you can specify project aware upstream flow dependencies
1523
- by specifying the fully qualified project_flow_name.
1524
- ```
1525
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1526
- ```
1527
- or
1528
- ```
1529
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1530
- ```
1531
-
1532
- You can also specify just the project or project branch (other values will be
1533
- inferred from the current project or project branch):
1534
- ```
1535
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1536
- ```
1537
-
1538
- Note that `branch` is typically one of:
1539
- - `prod`
1540
- - `user.bob`
1541
- - `test.my_experiment`
1542
- - `prod.staging`
1543
-
1544
- Parameters
1545
- ----------
1546
- flow : Union[str, Dict[str, str]], optional, default None
1547
- Upstream flow dependency for this flow.
1548
- flows : List[Union[str, Dict[str, str]]], default []
1549
- Upstream flow dependencies for this flow.
1516
+ events : List[Union[str, Dict[str, Any]]], default []
1517
+ Events dependency for this flow.
1550
1518
  options : Dict[str, Any], default {}
1551
1519
  Backend-specific configuration for tuning eventing behavior.
1552
1520
 
@@ -1555,51 +1523,47 @@ def trigger_on_finish(*, flow: typing.Union[str, typing.Dict[str, str], None] =
1555
1523
  ...
1556
1524
 
1557
1525
  @typing.overload
1558
- def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1526
+ def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1559
1527
  ...
1560
1528
 
1561
- def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, flow: typing.Union[str, typing.Dict[str, str], None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}):
1529
+ 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] = {}):
1562
1530
  """
1563
- Specifies the flow(s) that this flow depends on.
1531
+ Specifies the event(s) that this flow depends on.
1564
1532
 
1565
1533
  ```
1566
- @trigger_on_finish(flow='FooFlow')
1534
+ @trigger(event='foo')
1567
1535
  ```
1568
1536
  or
1569
1537
  ```
1570
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1538
+ @trigger(events=['foo', 'bar'])
1571
1539
  ```
1572
- This decorator respects the @project decorator and triggers the flow
1573
- when upstream runs within the same namespace complete successfully
1574
1540
 
1575
- Additionally, you can specify project aware upstream flow dependencies
1576
- by specifying the fully qualified project_flow_name.
1541
+ Additionally, you can specify the parameter mappings
1542
+ to map event payload to Metaflow parameters for the flow.
1577
1543
  ```
1578
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1544
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1579
1545
  ```
1580
1546
  or
1581
1547
  ```
1582
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1548
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1549
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1583
1550
  ```
1584
1551
 
1585
- You can also specify just the project or project branch (other values will be
1586
- inferred from the current project or project branch):
1552
+ 'parameters' can also be a list of strings and tuples like so:
1587
1553
  ```
1588
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1554
+ @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1555
+ ```
1556
+ This is equivalent to:
1557
+ ```
1558
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1589
1559
  ```
1590
-
1591
- Note that `branch` is typically one of:
1592
- - `prod`
1593
- - `user.bob`
1594
- - `test.my_experiment`
1595
- - `prod.staging`
1596
1560
 
1597
1561
  Parameters
1598
1562
  ----------
1599
- flow : Union[str, Dict[str, str]], optional, default None
1600
- Upstream flow dependency for this flow.
1601
- flows : List[Union[str, Dict[str, str]]], default []
1602
- Upstream flow dependencies for this flow.
1563
+ event : Union[str, Dict[str, Any]], optional, default None
1564
+ Event dependency for this flow.
1565
+ events : List[Union[str, Dict[str, Any]]], default []
1566
+ Events dependency for this flow.
1603
1567
  options : Dict[str, Any], default {}
1604
1568
  Backend-specific configuration for tuning eventing behavior.
1605
1569
 
@@ -1695,6 +1659,48 @@ def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packa
1695
1659
  """
1696
1660
  ...
1697
1661
 
1662
+ 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]]:
1663
+ """
1664
+ 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)
1665
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1666
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1667
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
1668
+ starts only after all sensors finish.
1669
+
1670
+ Parameters
1671
+ ----------
1672
+ timeout : int
1673
+ Time, in seconds before the task times out and fails. (Default: 3600)
1674
+ poke_interval : int
1675
+ Time in seconds that the job should wait in between each try. (Default: 60)
1676
+ mode : str
1677
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1678
+ exponential_backoff : bool
1679
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1680
+ pool : str
1681
+ the slot pool this task should run in,
1682
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1683
+ soft_fail : bool
1684
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1685
+ name : str
1686
+ Name of the sensor on Airflow
1687
+ description : str
1688
+ Description of sensor in the Airflow UI
1689
+ bucket_key : Union[str, List[str]]
1690
+ The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1691
+ When it's specified as a full s3:// url, please leave `bucket_name` as None
1692
+ bucket_name : str
1693
+ Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1694
+ When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1695
+ wildcard_match : bool
1696
+ whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1697
+ aws_conn_id : str
1698
+ a reference to the s3 connection on Airflow. (Default: None)
1699
+ verify : bool
1700
+ Whether or not to verify SSL certificates for S3 connection. (Default: None)
1701
+ """
1702
+ ...
1703
+
1698
1704
  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]]:
1699
1705
  """
1700
1706
  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.
@@ -1737,112 +1743,106 @@ def airflow_external_task_sensor(*, timeout: int, poke_interval: int, mode: str,
1737
1743
  """
1738
1744
  ...
1739
1745
 
1740
- def project(*, name: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1746
+ @typing.overload
1747
+ def trigger_on_finish(*, flow: typing.Union[str, typing.Dict[str, str], None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1741
1748
  """
1742
- Specifies what flows belong to the same project.
1749
+ Specifies the flow(s) that this flow depends on.
1743
1750
 
1744
- A project-specific namespace is created for all flows that
1745
- use the same `@project(name)`.
1751
+ ```
1752
+ @trigger_on_finish(flow='FooFlow')
1753
+ ```
1754
+ or
1755
+ ```
1756
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1757
+ ```
1758
+ This decorator respects the @project decorator and triggers the flow
1759
+ when upstream runs within the same namespace complete successfully
1746
1760
 
1747
- Parameters
1748
- ----------
1749
- name : str
1750
- Project name. Make sure that the name is unique amongst all
1751
- projects that use the same production scheduler. The name may
1752
- contain only lowercase alphanumeric characters and underscores.
1761
+ Additionally, you can specify project aware upstream flow dependencies
1762
+ by specifying the fully qualified project_flow_name.
1763
+ ```
1764
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1765
+ ```
1766
+ or
1767
+ ```
1768
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1769
+ ```
1753
1770
 
1771
+ You can also specify just the project or project branch (other values will be
1772
+ inferred from the current project or project branch):
1773
+ ```
1774
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1775
+ ```
1754
1776
 
1755
- """
1756
- ...
1757
-
1758
- @typing.overload
1759
- 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]]:
1760
- """
1761
- Specifies the times when the flow should be run when running on a
1762
- production scheduler.
1777
+ Note that `branch` is typically one of:
1778
+ - `prod`
1779
+ - `user.bob`
1780
+ - `test.my_experiment`
1781
+ - `prod.staging`
1763
1782
 
1764
1783
  Parameters
1765
1784
  ----------
1766
- hourly : bool, default False
1767
- Run the workflow hourly.
1768
- daily : bool, default True
1769
- Run the workflow daily.
1770
- weekly : bool, default False
1771
- Run the workflow weekly.
1772
- cron : str, optional, default None
1773
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1774
- specified by this expression.
1775
- timezone : str, optional, default None
1776
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1777
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1785
+ flow : Union[str, Dict[str, str]], optional, default None
1786
+ Upstream flow dependency for this flow.
1787
+ flows : List[Union[str, Dict[str, str]]], default []
1788
+ Upstream flow dependencies for this flow.
1789
+ options : Dict[str, Any], default {}
1790
+ Backend-specific configuration for tuning eventing behavior.
1791
+
1792
+
1778
1793
  """
1779
1794
  ...
1780
1795
 
1781
1796
  @typing.overload
1782
- def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1797
+ def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1783
1798
  ...
1784
1799
 
1785
- 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):
1800
+ def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, flow: typing.Union[str, typing.Dict[str, str], None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}):
1786
1801
  """
1787
- Specifies the times when the flow should be run when running on a
1788
- production scheduler.
1802
+ Specifies the flow(s) that this flow depends on.
1789
1803
 
1790
- Parameters
1791
- ----------
1792
- hourly : bool, default False
1793
- Run the workflow hourly.
1794
- daily : bool, default True
1795
- Run the workflow daily.
1796
- weekly : bool, default False
1797
- Run the workflow weekly.
1798
- cron : str, optional, default None
1799
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1800
- specified by this expression.
1801
- timezone : str, optional, default None
1802
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1803
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1804
- """
1805
- ...
1806
-
1807
- 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]]:
1808
- """
1809
- 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)
1810
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1811
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1812
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
1813
- starts only after all sensors finish.
1804
+ ```
1805
+ @trigger_on_finish(flow='FooFlow')
1806
+ ```
1807
+ or
1808
+ ```
1809
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1810
+ ```
1811
+ This decorator respects the @project decorator and triggers the flow
1812
+ when upstream runs within the same namespace complete successfully
1813
+
1814
+ Additionally, you can specify project aware upstream flow dependencies
1815
+ by specifying the fully qualified project_flow_name.
1816
+ ```
1817
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1818
+ ```
1819
+ or
1820
+ ```
1821
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1822
+ ```
1823
+
1824
+ You can also specify just the project or project branch (other values will be
1825
+ inferred from the current project or project branch):
1826
+ ```
1827
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1828
+ ```
1829
+
1830
+ Note that `branch` is typically one of:
1831
+ - `prod`
1832
+ - `user.bob`
1833
+ - `test.my_experiment`
1834
+ - `prod.staging`
1814
1835
 
1815
1836
  Parameters
1816
1837
  ----------
1817
- timeout : int
1818
- Time, in seconds before the task times out and fails. (Default: 3600)
1819
- poke_interval : int
1820
- Time in seconds that the job should wait in between each try. (Default: 60)
1821
- mode : str
1822
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1823
- exponential_backoff : bool
1824
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1825
- pool : str
1826
- the slot pool this task should run in,
1827
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1828
- soft_fail : bool
1829
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1830
- name : str
1831
- Name of the sensor on Airflow
1832
- description : str
1833
- Description of sensor in the Airflow UI
1834
- bucket_key : Union[str, List[str]]
1835
- The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1836
- When it's specified as a full s3:// url, please leave `bucket_name` as None
1837
- bucket_name : str
1838
- Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1839
- When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1840
- wildcard_match : bool
1841
- whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1842
- aws_conn_id : str
1843
- a reference to the s3 connection on Airflow. (Default: None)
1844
- verify : bool
1845
- Whether or not to verify SSL certificates for S3 connection. (Default: None)
1838
+ flow : Union[str, Dict[str, str]], optional, default None
1839
+ Upstream flow dependency for this flow.
1840
+ flows : List[Union[str, Dict[str, str]]], default []
1841
+ Upstream flow dependencies for this flow.
1842
+ options : Dict[str, Any], default {}
1843
+ Backend-specific configuration for tuning eventing behavior.
1844
+
1845
+
1846
1846
  """
1847
1847
  ...
1848
1848