ob-metaflow-stubs 4.3__py2.py3-none-any.whl → 4.4__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 (145) hide show
  1. metaflow-stubs/__init__.pyi +505 -474
  2. metaflow-stubs/cards.pyi +4 -4
  3. metaflow-stubs/cli.pyi +1 -1
  4. metaflow-stubs/client/__init__.pyi +1 -1
  5. metaflow-stubs/client/core.pyi +6 -6
  6. metaflow-stubs/client/filecache.pyi +1 -1
  7. metaflow-stubs/clone_util.pyi +1 -1
  8. metaflow-stubs/events.pyi +2 -2
  9. metaflow-stubs/exception.pyi +1 -1
  10. metaflow-stubs/flowspec.pyi +4 -4
  11. metaflow-stubs/generated_for.txt +1 -1
  12. metaflow-stubs/includefile.pyi +3 -3
  13. metaflow-stubs/metadata/metadata.pyi +1 -1
  14. metaflow-stubs/metadata/util.pyi +1 -1
  15. metaflow-stubs/metaflow_config.pyi +1 -1
  16. metaflow-stubs/metaflow_current.pyi +16 -16
  17. metaflow-stubs/mflog/mflog.pyi +1 -1
  18. metaflow-stubs/multicore_utils.pyi +1 -1
  19. metaflow-stubs/parameters.pyi +2 -2
  20. metaflow-stubs/plugins/__init__.pyi +1 -1
  21. metaflow-stubs/plugins/airflow/__init__.pyi +1 -1
  22. metaflow-stubs/plugins/airflow/airflow.pyi +2 -2
  23. metaflow-stubs/plugins/airflow/airflow_cli.pyi +2 -2
  24. metaflow-stubs/plugins/airflow/airflow_decorator.pyi +1 -1
  25. metaflow-stubs/plugins/airflow/airflow_utils.pyi +1 -1
  26. metaflow-stubs/plugins/airflow/exception.pyi +1 -1
  27. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +1 -1
  28. metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi +2 -2
  29. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +2 -2
  30. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +2 -2
  31. metaflow-stubs/plugins/argo/__init__.pyi +1 -1
  32. metaflow-stubs/plugins/argo/argo_client.pyi +1 -1
  33. metaflow-stubs/plugins/argo/argo_events.pyi +1 -1
  34. metaflow-stubs/plugins/argo/argo_workflows.pyi +3 -3
  35. metaflow-stubs/plugins/argo/argo_workflows_cli.pyi +3 -3
  36. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +2 -2
  37. metaflow-stubs/plugins/aws/__init__.pyi +1 -1
  38. metaflow-stubs/plugins/aws/aws_client.pyi +1 -1
  39. metaflow-stubs/plugins/aws/aws_utils.pyi +1 -1
  40. metaflow-stubs/plugins/aws/batch/__init__.pyi +1 -1
  41. metaflow-stubs/plugins/aws/batch/batch.pyi +1 -1
  42. metaflow-stubs/plugins/aws/batch/batch_cli.pyi +1 -1
  43. metaflow-stubs/plugins/aws/batch/batch_client.pyi +1 -1
  44. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +2 -2
  45. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +1 -1
  46. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +2 -2
  47. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +1 -1
  48. metaflow-stubs/plugins/aws/step_functions/dynamo_db_client.pyi +1 -1
  49. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +1 -1
  50. metaflow-stubs/plugins/aws/step_functions/production_token.pyi +1 -1
  51. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +1 -1
  52. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +1 -1
  53. metaflow-stubs/plugins/aws/step_functions/step_functions_cli.pyi +3 -3
  54. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +1 -1
  55. metaflow-stubs/plugins/aws/step_functions/step_functions_decorator.pyi +1 -1
  56. metaflow-stubs/plugins/azure/__init__.pyi +1 -1
  57. metaflow-stubs/plugins/azure/azure_credential.pyi +1 -1
  58. metaflow-stubs/plugins/azure/azure_exceptions.pyi +1 -1
  59. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +2 -2
  60. metaflow-stubs/plugins/azure/azure_utils.pyi +1 -1
  61. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +1 -1
  62. metaflow-stubs/plugins/azure/includefile_support.pyi +1 -1
  63. metaflow-stubs/plugins/cards/__init__.pyi +1 -1
  64. metaflow-stubs/plugins/cards/card_cli.pyi +3 -3
  65. metaflow-stubs/plugins/cards/card_client.pyi +2 -2
  66. metaflow-stubs/plugins/cards/card_creator.pyi +1 -1
  67. metaflow-stubs/plugins/cards/card_datastore.pyi +1 -1
  68. metaflow-stubs/plugins/cards/card_decorator.pyi +2 -2
  69. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +1 -1
  70. metaflow-stubs/plugins/cards/card_modules/basic.pyi +3 -3
  71. metaflow-stubs/plugins/cards/card_modules/card.pyi +1 -1
  72. metaflow-stubs/plugins/cards/card_modules/chevron/__init__.pyi +1 -1
  73. metaflow-stubs/plugins/cards/card_modules/chevron/main.pyi +1 -1
  74. metaflow-stubs/plugins/cards/card_modules/chevron/metadata.pyi +1 -1
  75. metaflow-stubs/plugins/cards/card_modules/chevron/renderer.pyi +1 -1
  76. metaflow-stubs/plugins/cards/card_modules/chevron/tokenizer.pyi +1 -1
  77. metaflow-stubs/plugins/cards/card_modules/components.pyi +3 -3
  78. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +1 -1
  79. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +1 -1
  80. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +2 -2
  81. metaflow-stubs/plugins/cards/card_resolver.pyi +1 -1
  82. metaflow-stubs/plugins/cards/component_serializer.pyi +3 -3
  83. metaflow-stubs/plugins/cards/exception.pyi +1 -1
  84. metaflow-stubs/plugins/catch_decorator.pyi +2 -2
  85. metaflow-stubs/plugins/datatools/__init__.pyi +1 -1
  86. metaflow-stubs/plugins/datatools/local.pyi +1 -1
  87. metaflow-stubs/plugins/datatools/s3/__init__.pyi +1 -1
  88. metaflow-stubs/plugins/datatools/s3/s3.pyi +5 -5
  89. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +1 -1
  90. metaflow-stubs/plugins/datatools/s3/s3util.pyi +1 -1
  91. metaflow-stubs/plugins/debug_logger.pyi +1 -1
  92. metaflow-stubs/plugins/debug_monitor.pyi +1 -1
  93. metaflow-stubs/plugins/environment_decorator.pyi +1 -1
  94. metaflow-stubs/plugins/events_decorator.pyi +2 -2
  95. metaflow-stubs/plugins/frameworks/__init__.pyi +1 -1
  96. metaflow-stubs/plugins/frameworks/pytorch.pyi +2 -2
  97. metaflow-stubs/plugins/gcp/__init__.pyi +1 -1
  98. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +2 -2
  99. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +1 -1
  100. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +1 -1
  101. metaflow-stubs/plugins/gcp/gs_utils.pyi +1 -1
  102. metaflow-stubs/plugins/gcp/includefile_support.pyi +1 -1
  103. metaflow-stubs/plugins/kubernetes/__init__.pyi +1 -1
  104. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +1 -1
  105. metaflow-stubs/plugins/kubernetes/kubernetes_cli.pyi +3 -3
  106. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +1 -1
  107. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +2 -2
  108. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +1 -1
  109. metaflow-stubs/plugins/logs_cli.pyi +2 -2
  110. metaflow-stubs/plugins/package_cli.pyi +1 -1
  111. metaflow-stubs/plugins/parallel_decorator.pyi +1 -1
  112. metaflow-stubs/plugins/perimeters.pyi +1 -1
  113. metaflow-stubs/plugins/project_decorator.pyi +2 -2
  114. metaflow-stubs/plugins/pypi/__init__.pyi +1 -1
  115. metaflow-stubs/plugins/pypi/conda_decorator.pyi +1 -1
  116. metaflow-stubs/plugins/pypi/conda_environment.pyi +1 -1
  117. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +1 -1
  118. metaflow-stubs/plugins/pypi/pypi_environment.pyi +1 -1
  119. metaflow-stubs/plugins/pypi/utils.pyi +1 -1
  120. metaflow-stubs/plugins/resources_decorator.pyi +1 -1
  121. metaflow-stubs/plugins/retry_decorator.pyi +1 -1
  122. metaflow-stubs/plugins/secrets/__init__.pyi +1 -1
  123. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +1 -1
  124. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +1 -1
  125. metaflow-stubs/plugins/storage_executor.pyi +1 -1
  126. metaflow-stubs/plugins/tag_cli.pyi +3 -3
  127. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +1 -1
  128. metaflow-stubs/plugins/timeout_decorator.pyi +2 -2
  129. metaflow-stubs/procpoll.pyi +1 -1
  130. metaflow-stubs/profilers/__init__.pyi +1 -1
  131. metaflow-stubs/pylint_wrapper.pyi +1 -1
  132. metaflow-stubs/runner/__init__.pyi +1 -1
  133. metaflow-stubs/runner/metaflow_runner.pyi +3 -3
  134. metaflow-stubs/runner/nbrun.pyi +1 -1
  135. metaflow-stubs/runner/subprocess_manager.pyi +1 -1
  136. metaflow-stubs/system/__init__.pyi +1 -1
  137. metaflow-stubs/system/system_logger.pyi +1 -1
  138. metaflow-stubs/system/system_monitor.pyi +1 -1
  139. metaflow-stubs/tagging_util.pyi +1 -1
  140. metaflow-stubs/tuple_util.pyi +1 -1
  141. {ob_metaflow_stubs-4.3.dist-info → ob_metaflow_stubs-4.4.dist-info}/METADATA +1 -1
  142. ob_metaflow_stubs-4.4.dist-info/RECORD +145 -0
  143. ob_metaflow_stubs-4.3.dist-info/RECORD +0 -145
  144. {ob_metaflow_stubs-4.3.dist-info → ob_metaflow_stubs-4.4.dist-info}/WHEEL +0 -0
  145. {ob_metaflow_stubs-4.3.dist-info → ob_metaflow_stubs-4.4.dist-info}/top_level.txt +0 -0
@@ -1,25 +1,25 @@
1
1
  ##################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
3
  # MF version: 2.12.7.1+ob(v1) #
4
- # Generated on 2024-07-09T16:19:31.537205 #
4
+ # Generated on 2024-07-11T17:01:16.236498 #
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.events
11
12
  import metaflow.client.core
12
- import metaflow.metaflow_current
13
- import io
14
- import metaflow.plugins.datatools.s3.s3
15
- import datetime
16
- import metaflow.datastore.inputs
13
+ import typing
17
14
  import metaflow.flowspec
15
+ import metaflow.plugins.datatools.s3.s3
18
16
  import metaflow.parameters
19
- import metaflow._vendor.click.types
20
- import metaflow.events
21
- import typing
22
17
  import metaflow.runner.metaflow_runner
18
+ import io
19
+ import metaflow._vendor.click.types
20
+ import datetime
21
+ import metaflow.datastore.inputs
22
+ import metaflow.metaflow_current
23
23
  FlowSpecDerived = typing.TypeVar("FlowSpecDerived", bound="FlowSpec", contravariant=False, covariant=False)
24
24
  StepFlag = typing.NewType("StepFlag", bool)
25
25
 
@@ -784,6 +784,96 @@ def conda(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
784
784
  """
785
785
  ...
786
786
 
787
+ @typing.overload
788
+ 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]]]:
789
+ """
790
+ Specifies a timeout for your step.
791
+
792
+ This decorator is useful if this step may hang indefinitely.
793
+
794
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
795
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
796
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
797
+
798
+ Note that all the values specified in parameters are added together so if you specify
799
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
800
+
801
+ Parameters
802
+ ----------
803
+ seconds : int, default 0
804
+ Number of seconds to wait prior to timing out.
805
+ minutes : int, default 0
806
+ Number of minutes to wait prior to timing out.
807
+ hours : int, default 0
808
+ Number of hours to wait prior to timing out.
809
+ """
810
+ ...
811
+
812
+ @typing.overload
813
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
814
+ ...
815
+
816
+ @typing.overload
817
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
818
+ ...
819
+
820
+ 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):
821
+ """
822
+ Specifies a timeout for your step.
823
+
824
+ This decorator is useful if this step may hang indefinitely.
825
+
826
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
827
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
828
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
829
+
830
+ Note that all the values specified in parameters are added together so if you specify
831
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
832
+
833
+ Parameters
834
+ ----------
835
+ seconds : int, default 0
836
+ Number of seconds to wait prior to timing out.
837
+ minutes : int, default 0
838
+ Number of minutes to wait prior to timing out.
839
+ hours : int, default 0
840
+ Number of hours to wait prior to timing out.
841
+ """
842
+ ...
843
+
844
+ @typing.overload
845
+ 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]]]:
846
+ """
847
+ Specifies secrets to be retrieved and injected as environment variables prior to
848
+ the execution of a step.
849
+
850
+ Parameters
851
+ ----------
852
+ sources : List[Union[str, Dict[str, Any]]], default: []
853
+ List of secret specs, defining how the secrets are to be retrieved
854
+ """
855
+ ...
856
+
857
+ @typing.overload
858
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
859
+ ...
860
+
861
+ @typing.overload
862
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
863
+ ...
864
+
865
+ 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]]] = []):
866
+ """
867
+ Specifies secrets to be retrieved and injected as environment variables prior to
868
+ the execution of a step.
869
+
870
+ Parameters
871
+ ----------
872
+ sources : List[Union[str, Dict[str, Any]]], default: []
873
+ List of secret specs, defining how the secrets are to be retrieved
874
+ """
875
+ ...
876
+
787
877
  @typing.overload
788
878
  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]]]:
789
879
  """
@@ -931,105 +1021,93 @@ def batch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
931
1021
  """
932
1022
  ...
933
1023
 
934
- @typing.overload
935
- 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]]]:
936
- """
937
- Specifies the number of times the task corresponding
938
- to a step needs to be retried.
939
-
940
- This decorator is useful for handling transient errors, such as networking issues.
941
- If your task contains operations that can't be retried safely, e.g. database updates,
942
- it is advisable to annotate it with `@retry(times=0)`.
943
-
944
- This can be used in conjunction with the `@catch` decorator. The `@catch`
945
- decorator will execute a no-op task after all retries have been exhausted,
946
- ensuring that the flow execution can continue.
947
-
948
- Parameters
949
- ----------
950
- times : int, default 3
951
- Number of times to retry this task.
952
- minutes_between_retries : int, default 2
953
- Number of minutes between retries.
954
- """
955
- ...
956
-
957
- @typing.overload
958
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
959
- ...
960
-
961
- @typing.overload
962
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
963
- ...
964
-
965
- 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):
1024
+ 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]]]:
966
1025
  """
967
- Specifies the number of times the task corresponding
968
- to a step needs to be retried.
969
-
970
- This decorator is useful for handling transient errors, such as networking issues.
971
- If your task contains operations that can't be retried safely, e.g. database updates,
972
- it is advisable to annotate it with `@retry(times=0)`.
973
-
974
- This can be used in conjunction with the `@catch` decorator. The `@catch`
975
- decorator will execute a no-op task after all retries have been exhausted,
976
- ensuring that the flow execution can continue.
1026
+ Specifies that this step should execute on Kubernetes.
977
1027
 
978
1028
  Parameters
979
1029
  ----------
980
- times : int, default 3
981
- Number of times to retry this task.
982
- minutes_between_retries : int, default 2
983
- Number of minutes between retries.
1030
+ cpu : int, default 1
1031
+ Number of CPUs required for this step. If `@resources` is
1032
+ also present, the maximum value from all decorators is used.
1033
+ memory : int, default 4096
1034
+ Memory size (in MB) required for this step. If
1035
+ `@resources` is also present, the maximum value from all decorators is
1036
+ used.
1037
+ disk : int, default 10240
1038
+ Disk size (in MB) required for this step. If
1039
+ `@resources` is also present, the maximum value from all decorators is
1040
+ used.
1041
+ image : str, optional, default None
1042
+ Docker image to use when launching on Kubernetes. If not specified, and
1043
+ METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
1044
+ not, a default Docker image mapping to the current version of Python is used.
1045
+ image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
1046
+ If given, the imagePullPolicy to be applied to the Docker image of the step.
1047
+ service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
1048
+ Kubernetes service account to use when launching pod in Kubernetes.
1049
+ secrets : List[str], optional, default None
1050
+ Kubernetes secrets to use when launching pod in Kubernetes. These
1051
+ secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
1052
+ in Metaflow configuration.
1053
+ namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
1054
+ Kubernetes namespace to use when launching pod in Kubernetes.
1055
+ gpu : int, optional, default None
1056
+ Number of GPUs required for this step. A value of zero implies that
1057
+ the scheduled node should not have GPUs.
1058
+ gpu_vendor : str, default KUBERNETES_GPU_VENDOR
1059
+ The vendor of the GPUs to be used for this step.
1060
+ tolerations : List[str], default []
1061
+ The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
1062
+ Kubernetes tolerations to use when launching pod in Kubernetes.
1063
+ use_tmpfs : bool, default False
1064
+ This enables an explicit tmpfs mount for this step.
1065
+ tmpfs_tempdir : bool, default True
1066
+ sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
1067
+ tmpfs_size : int, optional, default: None
1068
+ The value for the size (in MiB) of the tmpfs mount for this step.
1069
+ This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
1070
+ memory allocated for this step.
1071
+ tmpfs_path : str, optional, default /metaflow_temp
1072
+ Path to tmpfs mount for this step.
1073
+ persistent_volume_claims : Dict[str, str], optional, default None
1074
+ A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
1075
+ volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
1076
+ shared_memory: int, optional
1077
+ Shared memory size (in MiB) required for this step
1078
+ port: int, optional
1079
+ Port number to specify in the Kubernetes job object
984
1080
  """
985
1081
  ...
986
1082
 
987
1083
  @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]]]:
1084
+ 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]]]:
989
1085
  """
990
- Specifies the PyPI packages for the step.
991
-
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.
1086
+ Specifies environment variables to be set prior to the execution of a step.
996
1087
 
997
1088
  Parameters
998
1089
  ----------
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.
1090
+ vars : Dict[str, str], default {}
1091
+ Dictionary of environment variables to set.
1005
1092
  """
1006
1093
  ...
1007
1094
 
1008
1095
  @typing.overload
1009
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1096
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1010
1097
  ...
1011
1098
 
1012
1099
  @typing.overload
1013
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1100
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1014
1101
  ...
1015
1102
 
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):
1103
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
1017
1104
  """
1018
- Specifies the PyPI packages for the step.
1019
-
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.
1105
+ Specifies environment variables to be set prior to the execution of a step.
1024
1106
 
1025
1107
  Parameters
1026
1108
  ----------
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.
1109
+ vars : Dict[str, str], default {}
1110
+ Dictionary of environment variables to set.
1033
1111
  """
1034
1112
  ...
1035
1113
 
@@ -1111,90 +1189,51 @@ def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None]
1111
1189
  ...
1112
1190
 
1113
1191
  @typing.overload
1114
- 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]]]:
1192
+ 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]]]:
1115
1193
  """
1116
- Specifies environment variables to be set prior to the execution of a step.
1194
+ Specifies that the step will success under all circumstances.
1195
+
1196
+ The decorator will create an optional artifact, specified by `var`, which
1197
+ contains the exception raised. You can use it to detect the presence
1198
+ of errors, indicating that all happy-path artifacts produced by the step
1199
+ are missing.
1117
1200
 
1118
1201
  Parameters
1119
1202
  ----------
1120
- vars : Dict[str, str], default {}
1121
- Dictionary of environment variables to set.
1122
- """
1123
- ...
1124
-
1125
- @typing.overload
1126
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1127
- ...
1128
-
1129
- @typing.overload
1130
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1131
- ...
1132
-
1133
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
1134
- """
1135
- Specifies environment variables to be set prior to the execution of a step.
1136
-
1137
- Parameters
1138
- ----------
1139
- vars : Dict[str, str], default {}
1140
- Dictionary of environment variables to set.
1141
- """
1142
- ...
1143
-
1144
- @typing.overload
1145
- 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]]]:
1146
- """
1147
- Specifies a timeout for your step.
1148
-
1149
- This decorator is useful if this step may hang indefinitely.
1150
-
1151
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
1152
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
1153
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
1154
-
1155
- Note that all the values specified in parameters are added together so if you specify
1156
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
1157
-
1158
- Parameters
1159
- ----------
1160
- seconds : int, default 0
1161
- Number of seconds to wait prior to timing out.
1162
- minutes : int, default 0
1163
- Number of minutes to wait prior to timing out.
1164
- hours : int, default 0
1165
- Number of hours to wait prior to timing out.
1203
+ var : str, optional, default None
1204
+ Name of the artifact in which to store the caught exception.
1205
+ If not specified, the exception is not stored.
1206
+ print_exception : bool, default True
1207
+ Determines whether or not the exception is printed to
1208
+ stdout when caught.
1166
1209
  """
1167
1210
  ...
1168
1211
 
1169
1212
  @typing.overload
1170
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1213
+ def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1171
1214
  ...
1172
1215
 
1173
1216
  @typing.overload
1174
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1217
+ def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1175
1218
  ...
1176
1219
 
1177
- 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):
1220
+ 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):
1178
1221
  """
1179
- Specifies a timeout for your step.
1180
-
1181
- This decorator is useful if this step may hang indefinitely.
1182
-
1183
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
1184
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
1185
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
1222
+ Specifies that the step will success under all circumstances.
1186
1223
 
1187
- Note that all the values specified in parameters are added together so if you specify
1188
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
1224
+ The decorator will create an optional artifact, specified by `var`, which
1225
+ contains the exception raised. You can use it to detect the presence
1226
+ of errors, indicating that all happy-path artifacts produced by the step
1227
+ are missing.
1189
1228
 
1190
1229
  Parameters
1191
1230
  ----------
1192
- seconds : int, default 0
1193
- Number of seconds to wait prior to timing out.
1194
- minutes : int, default 0
1195
- Number of minutes to wait prior to timing out.
1196
- hours : int, default 0
1197
- Number of hours to wait prior to timing out.
1231
+ var : str, optional, default None
1232
+ Name of the artifact in which to store the caught exception.
1233
+ If not specified, the exception is not stored.
1234
+ print_exception : bool, default True
1235
+ Determines whether or not the exception is printed to
1236
+ stdout when caught.
1198
1237
  """
1199
1238
  ...
1200
1239
 
@@ -1250,397 +1289,294 @@ def card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typ
1250
1289
  ...
1251
1290
 
1252
1291
  @typing.overload
1253
- 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]]]:
1292
+ 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]]]:
1254
1293
  """
1255
- Specifies secrets to be retrieved and injected as environment variables prior to
1256
- the execution of a step.
1294
+ Specifies the number of times the task corresponding
1295
+ to a step needs to be retried.
1296
+
1297
+ This decorator is useful for handling transient errors, such as networking issues.
1298
+ If your task contains operations that can't be retried safely, e.g. database updates,
1299
+ it is advisable to annotate it with `@retry(times=0)`.
1300
+
1301
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
1302
+ decorator will execute a no-op task after all retries have been exhausted,
1303
+ ensuring that the flow execution can continue.
1257
1304
 
1258
1305
  Parameters
1259
1306
  ----------
1260
- sources : List[Union[str, Dict[str, Any]]], default: []
1261
- List of secret specs, defining how the secrets are to be retrieved
1307
+ times : int, default 3
1308
+ Number of times to retry this task.
1309
+ minutes_between_retries : int, default 2
1310
+ Number of minutes between retries.
1262
1311
  """
1263
1312
  ...
1264
1313
 
1265
1314
  @typing.overload
1266
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1315
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1267
1316
  ...
1268
1317
 
1269
1318
  @typing.overload
1270
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1319
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1271
1320
  ...
1272
1321
 
1273
- 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]]] = []):
1322
+ 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):
1274
1323
  """
1275
- Specifies secrets to be retrieved and injected as environment variables prior to
1276
- the execution of a step.
1324
+ Specifies the number of times the task corresponding
1325
+ to a step needs to be retried.
1277
1326
 
1278
- Parameters
1279
- ----------
1280
- sources : List[Union[str, Dict[str, Any]]], default: []
1281
- List of secret specs, defining how the secrets are to be retrieved
1282
- """
1283
- ...
1284
-
1285
- 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]]]:
1286
- """
1287
- Specifies that this step should execute on Kubernetes.
1327
+ This decorator is useful for handling transient errors, such as networking issues.
1328
+ If your task contains operations that can't be retried safely, e.g. database updates,
1329
+ it is advisable to annotate it with `@retry(times=0)`.
1330
+
1331
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
1332
+ decorator will execute a no-op task after all retries have been exhausted,
1333
+ ensuring that the flow execution can continue.
1288
1334
 
1289
1335
  Parameters
1290
1336
  ----------
1291
- cpu : int, default 1
1292
- Number of CPUs required for this step. If `@resources` is
1293
- also present, the maximum value from all decorators is used.
1294
- memory : int, default 4096
1295
- Memory size (in MB) required for this step. If
1296
- `@resources` is also present, the maximum value from all decorators is
1297
- used.
1298
- disk : int, default 10240
1299
- Disk size (in MB) required for this step. If
1300
- `@resources` is also present, the maximum value from all decorators is
1301
- used.
1302
- image : str, optional, default None
1303
- Docker image to use when launching on Kubernetes. If not specified, and
1304
- METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
1305
- not, a default Docker image mapping to the current version of Python is used.
1306
- image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
1307
- If given, the imagePullPolicy to be applied to the Docker image of the step.
1308
- service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
1309
- Kubernetes service account to use when launching pod in Kubernetes.
1310
- secrets : List[str], optional, default None
1311
- Kubernetes secrets to use when launching pod in Kubernetes. These
1312
- secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
1313
- in Metaflow configuration.
1314
- namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
1315
- Kubernetes namespace to use when launching pod in Kubernetes.
1316
- gpu : int, optional, default None
1317
- Number of GPUs required for this step. A value of zero implies that
1318
- the scheduled node should not have GPUs.
1319
- gpu_vendor : str, default KUBERNETES_GPU_VENDOR
1320
- The vendor of the GPUs to be used for this step.
1321
- tolerations : List[str], default []
1322
- The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
1323
- Kubernetes tolerations to use when launching pod in Kubernetes.
1324
- use_tmpfs : bool, default False
1325
- This enables an explicit tmpfs mount for this step.
1326
- tmpfs_tempdir : bool, default True
1327
- sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
1328
- tmpfs_size : int, optional, default: None
1329
- The value for the size (in MiB) of the tmpfs mount for this step.
1330
- This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
1331
- memory allocated for this step.
1332
- tmpfs_path : str, optional, default /metaflow_temp
1333
- Path to tmpfs mount for this step.
1334
- persistent_volume_claims : Dict[str, str], optional, default None
1335
- A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
1336
- volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
1337
- shared_memory: int, optional
1338
- Shared memory size (in MiB) required for this step
1339
- port: int, optional
1340
- Port number to specify in the Kubernetes job object
1337
+ times : int, default 3
1338
+ Number of times to retry this task.
1339
+ minutes_between_retries : int, default 2
1340
+ Number of minutes between retries.
1341
1341
  """
1342
1342
  ...
1343
1343
 
1344
1344
  @typing.overload
1345
- 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]]]:
1345
+ def pypi(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
1346
1346
  """
1347
- Specifies that the step will success under all circumstances.
1347
+ Specifies the PyPI packages for the step.
1348
1348
 
1349
- The decorator will create an optional artifact, specified by `var`, which
1350
- contains the exception raised. You can use it to detect the presence
1351
- of errors, indicating that all happy-path artifacts produced by the step
1352
- are missing.
1349
+ Information in this decorator will augment any
1350
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
1351
+ you can use `@pypi_base` to set packages required by all
1352
+ steps and use `@pypi` to specify step-specific overrides.
1353
1353
 
1354
1354
  Parameters
1355
1355
  ----------
1356
- var : str, optional, default None
1357
- Name of the artifact in which to store the caught exception.
1358
- If not specified, the exception is not stored.
1359
- print_exception : bool, default True
1360
- Determines whether or not the exception is printed to
1361
- stdout when caught.
1356
+ packages : Dict[str, str], default: {}
1357
+ Packages to use for this step. The key is the name of the package
1358
+ and the value is the version to use.
1359
+ python : str, optional, default: None
1360
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1361
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1362
1362
  """
1363
1363
  ...
1364
1364
 
1365
1365
  @typing.overload
1366
- def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1366
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1367
1367
  ...
1368
1368
 
1369
1369
  @typing.overload
1370
- def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1370
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1371
1371
  ...
1372
1372
 
1373
- 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):
1373
+ def pypi(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1374
1374
  """
1375
- Specifies that the step will success under all circumstances.
1375
+ Specifies the PyPI packages for the step.
1376
1376
 
1377
- The decorator will create an optional artifact, specified by `var`, which
1378
- contains the exception raised. You can use it to detect the presence
1379
- of errors, indicating that all happy-path artifacts produced by the step
1380
- are missing.
1377
+ Information in this decorator will augment any
1378
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
1379
+ you can use `@pypi_base` to set packages required by all
1380
+ steps and use `@pypi` to specify step-specific overrides.
1381
1381
 
1382
1382
  Parameters
1383
1383
  ----------
1384
- var : str, optional, default None
1385
- Name of the artifact in which to store the caught exception.
1386
- If not specified, the exception is not stored.
1387
- print_exception : bool, default True
1388
- Determines whether or not the exception is printed to
1389
- stdout when caught.
1384
+ packages : Dict[str, str], default: {}
1385
+ Packages to use for this step. The key is the name of the package
1386
+ and the value is the version to use.
1387
+ python : str, optional, default: None
1388
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1389
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1390
1390
  """
1391
1391
  ...
1392
1392
 
1393
1393
  @typing.overload
1394
- def conda_base(*, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1394
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1395
1395
  """
1396
- Specifies the Conda environment for all steps of the flow.
1397
-
1398
- Use `@conda_base` to set common libraries required by all
1399
- steps and use `@conda` to specify step-specific additions.
1396
+ Specifies the PyPI packages for all steps of the flow.
1400
1397
 
1398
+ Use `@pypi_base` to set common packages required by all
1399
+ steps and use `@pypi` to specify step-specific overrides.
1401
1400
  Parameters
1402
1401
  ----------
1403
- packages : Dict[str, str], default {}
1402
+ packages : Dict[str, str], default: {}
1404
1403
  Packages to use for this flow. The key is the name of the package
1405
1404
  and the value is the version to use.
1406
- libraries : Dict[str, str], default {}
1407
- Supported for backward compatibility. When used with packages, packages will take precedence.
1408
- python : str, optional, default None
1405
+ python : str, optional, default: None
1409
1406
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
1410
1407
  that the version used will correspond to the version of the Python interpreter used to start the run.
1411
- disabled : bool, default False
1412
- If set to True, disables Conda.
1413
1408
  """
1414
1409
  ...
1415
1410
 
1416
1411
  @typing.overload
1417
- def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1412
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1418
1413
  ...
1419
1414
 
1420
- def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
1415
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1421
1416
  """
1422
- Specifies the Conda environment for all steps of the flow.
1423
-
1424
- Use `@conda_base` to set common libraries required by all
1425
- steps and use `@conda` to specify step-specific additions.
1417
+ Specifies the PyPI packages for all steps of the flow.
1426
1418
 
1419
+ Use `@pypi_base` to set common packages required by all
1420
+ steps and use `@pypi` to specify step-specific overrides.
1427
1421
  Parameters
1428
1422
  ----------
1429
- packages : Dict[str, str], default {}
1423
+ packages : Dict[str, str], default: {}
1430
1424
  Packages to use for this flow. The key is the name of the package
1431
1425
  and the value is the version to use.
1432
- libraries : Dict[str, str], default {}
1433
- Supported for backward compatibility. When used with packages, packages will take precedence.
1434
- python : str, optional, default None
1426
+ python : str, optional, default: None
1435
1427
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
1436
1428
  that the version used will correspond to the version of the Python interpreter used to start the run.
1437
- disabled : bool, default False
1438
- If set to True, disables Conda.
1439
1429
  """
1440
1430
  ...
1441
1431
 
1442
- 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]]:
1432
+ def project(*, name: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1443
1433
  """
1444
- 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)
1445
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1446
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1447
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
1448
- starts only after all sensors finish.
1434
+ Specifies what flows belong to the same project.
1449
1435
 
1450
- Parameters
1451
- ----------
1452
- timeout : int
1453
- Time, in seconds before the task times out and fails. (Default: 3600)
1454
- poke_interval : int
1455
- Time in seconds that the job should wait in between each try. (Default: 60)
1456
- mode : str
1457
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1458
- exponential_backoff : bool
1459
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1460
- pool : str
1461
- the slot pool this task should run in,
1462
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1463
- soft_fail : bool
1464
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1465
- name : str
1466
- Name of the sensor on Airflow
1467
- description : str
1468
- Description of sensor in the Airflow UI
1469
- bucket_key : Union[str, List[str]]
1470
- The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1471
- When it's specified as a full s3:// url, please leave `bucket_name` as None
1472
- bucket_name : str
1473
- Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1474
- When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1475
- wildcard_match : bool
1476
- whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1477
- aws_conn_id : str
1478
- a reference to the s3 connection on Airflow. (Default: None)
1479
- verify : bool
1480
- Whether or not to verify SSL certificates for S3 connection. (Default: None)
1481
- """
1482
- ...
1483
-
1484
- 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]]:
1485
- """
1486
- 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.
1487
- This decorator only works when a flow is scheduled on Airflow and is compiled using `airflow create`. More than one `@airflow_external_task_sensor` can be added as a flow decorators. Adding more than one decorator will ensure that `start` step starts only after all sensors finish.
1436
+ A project-specific namespace is created for all flows that
1437
+ use the same `@project(name)`.
1488
1438
 
1489
1439
  Parameters
1490
1440
  ----------
1491
- timeout : int
1492
- Time, in seconds before the task times out and fails. (Default: 3600)
1493
- poke_interval : int
1494
- Time in seconds that the job should wait in between each try. (Default: 60)
1495
- mode : str
1496
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1497
- exponential_backoff : bool
1498
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1499
- pool : str
1500
- the slot pool this task should run in,
1501
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1502
- soft_fail : bool
1503
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1504
1441
  name : str
1505
- Name of the sensor on Airflow
1506
- description : str
1507
- Description of sensor in the Airflow UI
1508
- external_dag_id : str
1509
- The dag_id that contains the task you want to wait for.
1510
- external_task_ids : List[str]
1511
- The list of task_ids that you want to wait for.
1512
- If None (default value) the sensor waits for the DAG. (Default: None)
1513
- allowed_states : List[str]
1514
- Iterable of allowed states, (Default: ['success'])
1515
- failed_states : List[str]
1516
- Iterable of failed or dis-allowed states. (Default: None)
1517
- execution_delta : datetime.timedelta
1518
- time difference with the previous execution to look at,
1519
- the default is the same logical date as the current task or DAG. (Default: None)
1520
- check_existence: bool
1521
- Set to True to check if the external task exists or check if
1522
- the DAG to wait for exists. (Default: True)
1442
+ Project name. Make sure that the name is unique amongst all
1443
+ projects that use the same production scheduler. The name may
1444
+ contain only lowercase alphanumeric characters and underscores.
1445
+
1446
+
1523
1447
  """
1524
1448
  ...
1525
1449
 
1526
1450
  @typing.overload
1527
- 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]]:
1451
+ def conda_base(*, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1528
1452
  """
1529
- Specifies the flow(s) that this flow depends on.
1530
-
1531
- ```
1532
- @trigger_on_finish(flow='FooFlow')
1533
- ```
1534
- or
1535
- ```
1536
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1537
- ```
1538
- This decorator respects the @project decorator and triggers the flow
1539
- when upstream runs within the same namespace complete successfully
1540
-
1541
- Additionally, you can specify project aware upstream flow dependencies
1542
- by specifying the fully qualified project_flow_name.
1543
- ```
1544
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1545
- ```
1546
- or
1547
- ```
1548
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1549
- ```
1550
-
1551
- You can also specify just the project or project branch (other values will be
1552
- inferred from the current project or project branch):
1553
- ```
1554
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1555
- ```
1453
+ Specifies the Conda environment for all steps of the flow.
1556
1454
 
1557
- Note that `branch` is typically one of:
1558
- - `prod`
1559
- - `user.bob`
1560
- - `test.my_experiment`
1561
- - `prod.staging`
1455
+ Use `@conda_base` to set common libraries required by all
1456
+ steps and use `@conda` to specify step-specific additions.
1562
1457
 
1563
1458
  Parameters
1564
1459
  ----------
1565
- flow : Union[str, Dict[str, str]], optional, default None
1566
- Upstream flow dependency for this flow.
1567
- flows : List[Union[str, Dict[str, str]]], default []
1568
- Upstream flow dependencies for this flow.
1569
- options : Dict[str, Any], default {}
1570
- Backend-specific configuration for tuning eventing behavior.
1571
-
1572
-
1460
+ packages : Dict[str, str], default {}
1461
+ Packages to use for this flow. The key is the name of the package
1462
+ and the value is the version to use.
1463
+ libraries : Dict[str, str], default {}
1464
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1465
+ python : str, optional, default None
1466
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1467
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1468
+ disabled : bool, default False
1469
+ If set to True, disables Conda.
1573
1470
  """
1574
1471
  ...
1575
1472
 
1576
1473
  @typing.overload
1577
- def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1474
+ def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1578
1475
  ...
1579
1476
 
1580
- 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] = {}):
1477
+ def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
1581
1478
  """
1582
- Specifies the flow(s) that this flow depends on.
1583
-
1584
- ```
1585
- @trigger_on_finish(flow='FooFlow')
1586
- ```
1587
- or
1588
- ```
1589
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1590
- ```
1591
- This decorator respects the @project decorator and triggers the flow
1592
- when upstream runs within the same namespace complete successfully
1593
-
1594
- Additionally, you can specify project aware upstream flow dependencies
1595
- by specifying the fully qualified project_flow_name.
1596
- ```
1597
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1598
- ```
1599
- or
1600
- ```
1601
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1602
- ```
1603
-
1604
- You can also specify just the project or project branch (other values will be
1605
- inferred from the current project or project branch):
1606
- ```
1607
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1608
- ```
1479
+ Specifies the Conda environment for all steps of the flow.
1609
1480
 
1610
- Note that `branch` is typically one of:
1611
- - `prod`
1612
- - `user.bob`
1613
- - `test.my_experiment`
1614
- - `prod.staging`
1481
+ Use `@conda_base` to set common libraries required by all
1482
+ steps and use `@conda` to specify step-specific additions.
1615
1483
 
1616
1484
  Parameters
1617
1485
  ----------
1618
- flow : Union[str, Dict[str, str]], optional, default None
1619
- Upstream flow dependency for this flow.
1620
- flows : List[Union[str, Dict[str, str]]], default []
1621
- Upstream flow dependencies for this flow.
1622
- options : Dict[str, Any], default {}
1623
- Backend-specific configuration for tuning eventing behavior.
1624
-
1625
-
1486
+ packages : Dict[str, str], default {}
1487
+ Packages to use for this flow. The key is the name of the package
1488
+ and the value is the version to use.
1489
+ libraries : Dict[str, str], default {}
1490
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1491
+ python : str, optional, default None
1492
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1493
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1494
+ disabled : bool, default False
1495
+ If set to True, disables Conda.
1626
1496
  """
1627
1497
  ...
1628
1498
 
1629
- def project(*, name: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1499
+ def airflow_external_task_sensor(*, timeout: int, poke_interval: int, mode: str, exponential_backoff: bool, pool: str, soft_fail: bool, name: str, description: str, external_dag_id: str, external_task_ids: typing.List[str], allowed_states: typing.List[str], failed_states: typing.List[str], execution_delta: "datetime.timedelta", check_existence: bool) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1630
1500
  """
1631
- Specifies what flows belong to the same project.
1632
-
1633
- A project-specific namespace is created for all flows that
1634
- use the same `@project(name)`.
1501
+ The `@airflow_external_task_sensor` decorator attaches a Airflow [ExternalTaskSensor](https://airflow.apache.org/docs/apache-airflow/stable/_api/airflow/sensors/external_task/index.html#airflow.sensors.external_task.ExternalTaskSensor) before the start step of the flow.
1502
+ This decorator only works when a flow is scheduled on Airflow and is compiled using `airflow create`. More than one `@airflow_external_task_sensor` can be added as a flow decorators. Adding more than one decorator will ensure that `start` step starts only after all sensors finish.
1635
1503
 
1636
1504
  Parameters
1637
1505
  ----------
1506
+ timeout : int
1507
+ Time, in seconds before the task times out and fails. (Default: 3600)
1508
+ poke_interval : int
1509
+ Time in seconds that the job should wait in between each try. (Default: 60)
1510
+ mode : str
1511
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1512
+ exponential_backoff : bool
1513
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1514
+ pool : str
1515
+ the slot pool this task should run in,
1516
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1517
+ soft_fail : bool
1518
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1638
1519
  name : str
1639
- Project name. Make sure that the name is unique amongst all
1640
- projects that use the same production scheduler. The name may
1641
- contain only lowercase alphanumeric characters and underscores.
1642
-
1520
+ Name of the sensor on Airflow
1521
+ description : str
1522
+ Description of sensor in the Airflow UI
1523
+ external_dag_id : str
1524
+ The dag_id that contains the task you want to wait for.
1525
+ external_task_ids : List[str]
1526
+ The list of task_ids that you want to wait for.
1527
+ If None (default value) the sensor waits for the DAG. (Default: None)
1528
+ allowed_states : List[str]
1529
+ Iterable of allowed states, (Default: ['success'])
1530
+ failed_states : List[str]
1531
+ Iterable of failed or dis-allowed states. (Default: None)
1532
+ execution_delta : datetime.timedelta
1533
+ time difference with the previous execution to look at,
1534
+ the default is the same logical date as the current task or DAG. (Default: None)
1535
+ check_existence: bool
1536
+ Set to True to check if the external task exists or check if
1537
+ the DAG to wait for exists. (Default: True)
1538
+ """
1539
+ ...
1540
+
1541
+ 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]]:
1542
+ """
1543
+ 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)
1544
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1545
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1546
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
1547
+ starts only after all sensors finish.
1643
1548
 
1549
+ Parameters
1550
+ ----------
1551
+ timeout : int
1552
+ Time, in seconds before the task times out and fails. (Default: 3600)
1553
+ poke_interval : int
1554
+ Time in seconds that the job should wait in between each try. (Default: 60)
1555
+ mode : str
1556
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1557
+ exponential_backoff : bool
1558
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1559
+ pool : str
1560
+ the slot pool this task should run in,
1561
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1562
+ soft_fail : bool
1563
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1564
+ name : str
1565
+ Name of the sensor on Airflow
1566
+ description : str
1567
+ Description of sensor in the Airflow UI
1568
+ bucket_key : Union[str, List[str]]
1569
+ The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1570
+ When it's specified as a full s3:// url, please leave `bucket_name` as None
1571
+ bucket_name : str
1572
+ Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1573
+ When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1574
+ wildcard_match : bool
1575
+ whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1576
+ aws_conn_id : str
1577
+ a reference to the s3 connection on Airflow. (Default: None)
1578
+ verify : bool
1579
+ Whether or not to verify SSL certificates for S3 connection. (Default: None)
1644
1580
  """
1645
1581
  ...
1646
1582
 
@@ -1739,45 +1675,6 @@ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: t
1739
1675
  """
1740
1676
  ...
1741
1677
 
1742
- @typing.overload
1743
- def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1744
- """
1745
- Specifies the PyPI packages for all steps of the flow.
1746
-
1747
- Use `@pypi_base` to set common packages required by all
1748
- steps and use `@pypi` to specify step-specific overrides.
1749
- Parameters
1750
- ----------
1751
- packages : Dict[str, str], default: {}
1752
- Packages to use for this flow. The key is the name of the package
1753
- and the value is the version to use.
1754
- python : str, optional, default: None
1755
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1756
- that the version used will correspond to the version of the Python interpreter used to start the run.
1757
- """
1758
- ...
1759
-
1760
- @typing.overload
1761
- def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1762
- ...
1763
-
1764
- def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1765
- """
1766
- Specifies the PyPI packages for all steps of the flow.
1767
-
1768
- Use `@pypi_base` to set common packages required by all
1769
- steps and use `@pypi` to specify step-specific overrides.
1770
- Parameters
1771
- ----------
1772
- packages : Dict[str, str], default: {}
1773
- Packages to use for this flow. The key is the name of the package
1774
- and the value is the version to use.
1775
- python : str, optional, default: None
1776
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1777
- that the version used will correspond to the version of the Python interpreter used to start the run.
1778
- """
1779
- ...
1780
-
1781
1678
  @typing.overload
1782
1679
  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]]:
1783
1680
  """
@@ -1827,6 +1724,140 @@ def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly:
1827
1724
  """
1828
1725
  ...
1829
1726
 
1727
+ def nim(*, models: "list[NIM]", backend: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1728
+ """
1729
+ This decorator is used to run NIM containers in Metaflow tasks as sidecars.
1730
+
1731
+ User code call
1732
+ -----------
1733
+ @nim(
1734
+ models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
1735
+ backend='managed'
1736
+ )
1737
+
1738
+ Valid backend options
1739
+ ---------------------
1740
+ - 'managed': Outerbounds selects a compute provider based on the model.
1741
+ - 🚧 'dataplane': Run in your account.
1742
+
1743
+ Valid model options
1744
+ ----------------
1745
+ - 'meta/llama3-8b-instruct': 8B parameter model
1746
+ - 'meta/llama3-70b-instruct': 70B parameter model
1747
+ - Upon request, any model here: https://nvcf.ngc.nvidia.com/functions?filter=nvidia-functions
1748
+
1749
+ Parameters
1750
+ ----------
1751
+ models: list[NIM]
1752
+ List of NIM containers running models in sidecars.
1753
+ backend: str
1754
+ Compute provider to run the NIM container.
1755
+ """
1756
+ ...
1757
+
1758
+ @typing.overload
1759
+ 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]]:
1760
+ """
1761
+ Specifies the flow(s) that this flow depends on.
1762
+
1763
+ ```
1764
+ @trigger_on_finish(flow='FooFlow')
1765
+ ```
1766
+ or
1767
+ ```
1768
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1769
+ ```
1770
+ This decorator respects the @project decorator and triggers the flow
1771
+ when upstream runs within the same namespace complete successfully
1772
+
1773
+ Additionally, you can specify project aware upstream flow dependencies
1774
+ by specifying the fully qualified project_flow_name.
1775
+ ```
1776
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1777
+ ```
1778
+ or
1779
+ ```
1780
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1781
+ ```
1782
+
1783
+ You can also specify just the project or project branch (other values will be
1784
+ inferred from the current project or project branch):
1785
+ ```
1786
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1787
+ ```
1788
+
1789
+ Note that `branch` is typically one of:
1790
+ - `prod`
1791
+ - `user.bob`
1792
+ - `test.my_experiment`
1793
+ - `prod.staging`
1794
+
1795
+ Parameters
1796
+ ----------
1797
+ flow : Union[str, Dict[str, str]], optional, default None
1798
+ Upstream flow dependency for this flow.
1799
+ flows : List[Union[str, Dict[str, str]]], default []
1800
+ Upstream flow dependencies for this flow.
1801
+ options : Dict[str, Any], default {}
1802
+ Backend-specific configuration for tuning eventing behavior.
1803
+
1804
+
1805
+ """
1806
+ ...
1807
+
1808
+ @typing.overload
1809
+ def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1810
+ ...
1811
+
1812
+ 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] = {}):
1813
+ """
1814
+ Specifies the flow(s) that this flow depends on.
1815
+
1816
+ ```
1817
+ @trigger_on_finish(flow='FooFlow')
1818
+ ```
1819
+ or
1820
+ ```
1821
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1822
+ ```
1823
+ This decorator respects the @project decorator and triggers the flow
1824
+ when upstream runs within the same namespace complete successfully
1825
+
1826
+ Additionally, you can specify project aware upstream flow dependencies
1827
+ by specifying the fully qualified project_flow_name.
1828
+ ```
1829
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1830
+ ```
1831
+ or
1832
+ ```
1833
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1834
+ ```
1835
+
1836
+ You can also specify just the project or project branch (other values will be
1837
+ inferred from the current project or project branch):
1838
+ ```
1839
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1840
+ ```
1841
+
1842
+ Note that `branch` is typically one of:
1843
+ - `prod`
1844
+ - `user.bob`
1845
+ - `test.my_experiment`
1846
+ - `prod.staging`
1847
+
1848
+ Parameters
1849
+ ----------
1850
+ flow : Union[str, Dict[str, str]], optional, default None
1851
+ Upstream flow dependency for this flow.
1852
+ flows : List[Union[str, Dict[str, str]]], default []
1853
+ Upstream flow dependencies for this flow.
1854
+ options : Dict[str, Any], default {}
1855
+ Backend-specific configuration for tuning eventing behavior.
1856
+
1857
+
1858
+ """
1859
+ ...
1860
+
1830
1861
  def namespace(ns: typing.Optional[str]) -> typing.Optional[str]:
1831
1862
  """
1832
1863
  Switch namespace to the one provided.