metaflow-stubs 2.12.8__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 +590 -541
  2. metaflow-stubs/cards.pyi +5 -5
  3. metaflow-stubs/cli.pyi +2 -2
  4. metaflow-stubs/client/__init__.pyi +3 -3
  5. metaflow-stubs/client/core.pyi +6 -6
  6. metaflow-stubs/client/filecache.pyi +2 -2
  7. metaflow-stubs/clone_util.pyi +2 -2
  8. metaflow-stubs/events.pyi +2 -2
  9. metaflow-stubs/exception.pyi +2 -2
  10. metaflow-stubs/flowspec.pyi +6 -6
  11. metaflow-stubs/generated_for.txt +1 -1
  12. metaflow-stubs/includefile.pyi +4 -4
  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 +34 -11
  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 +6 -2
  21. metaflow-stubs/plugins/airflow/__init__.pyi +2 -2
  22. metaflow-stubs/plugins/airflow/airflow.pyi +5 -5
  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 +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 +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 +48 -4
  35. metaflow-stubs/plugins/argo/argo_workflows_cli.pyi +270 -5
  36. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +229 -4
  37. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +292 -0
  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 +24 -4
  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 +260 -0
  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 +3 -3
  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 +4 -4
  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 +4 -4
  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 +4 -4
  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 +6 -6
  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 +3 -3
  97. metaflow-stubs/plugins/frameworks/__init__.pyi +2 -2
  98. metaflow-stubs/plugins/frameworks/pytorch.pyi +5 -3
  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 +18 -10
  107. metaflow-stubs/plugins/kubernetes/kubernetes_cli.pyi +12 -12
  108. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +26 -3
  109. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +14 -31
  110. metaflow-stubs/plugins/kubernetes/kubernetes_job.pyi +26 -3
  111. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +82 -28
  112. metaflow-stubs/plugins/logs_cli.pyi +2 -2
  113. metaflow-stubs/plugins/package_cli.pyi +2 -2
  114. metaflow-stubs/plugins/parallel_decorator.pyi +47 -2
  115. metaflow-stubs/plugins/project_decorator.pyi +3 -3
  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 +4 -4
  119. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +2 -2
  120. metaflow-stubs/plugins/pypi/pypi_environment.pyi +3 -3
  121. metaflow-stubs/plugins/pypi/utils.pyi +2 -2
  122. metaflow-stubs/plugins/resources_decorator.pyi +2 -2
  123. metaflow-stubs/plugins/retry_decorator.pyi +2 -2
  124. metaflow-stubs/plugins/secrets/__init__.pyi +2 -2
  125. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +3 -3
  126. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +3 -3
  127. metaflow-stubs/plugins/storage_executor.pyi +2 -2
  128. metaflow-stubs/plugins/tag_cli.pyi +5 -5
  129. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +25 -2
  130. metaflow-stubs/plugins/timeout_decorator.pyi +2 -2
  131. metaflow-stubs/procpoll.pyi +2 -2
  132. metaflow-stubs/pylint_wrapper.pyi +2 -2
  133. metaflow-stubs/runner/__init__.pyi +2 -2
  134. metaflow-stubs/runner/deployer.pyi +354 -0
  135. metaflow-stubs/runner/metaflow_runner.pyi +11 -11
  136. metaflow-stubs/runner/nbdeploy.pyi +54 -0
  137. metaflow-stubs/runner/nbrun.pyi +7 -7
  138. metaflow-stubs/runner/subprocess_manager.pyi +2 -2
  139. metaflow-stubs/runner/utils.pyi +25 -0
  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.8.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.8.dist-info → metaflow_stubs-2.12.10.dist-info}/WHEEL +1 -1
  149. metaflow_stubs-2.12.8.dist-info/RECORD +0 -145
  150. {metaflow_stubs-2.12.8.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.8 #
4
- # Generated on 2024-07-16T15:51:55.346854 #
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.plugins.datatools.s3.s3
12
- import metaflow.runner.metaflow_runner
13
- import metaflow.datastore.inputs
14
- import metaflow.flowspec
15
- import metaflow._vendor.click.types
16
- import datetime
17
11
  import metaflow.metaflow_current
18
- import metaflow.events
19
- import typing
12
+ import metaflow._vendor.click.types
20
13
  import metaflow.client.core
14
+ import metaflow.events
21
15
  import metaflow.parameters
16
+ import datetime
17
+ import typing
18
+ import metaflow.plugins.datatools.s3.s3
19
+ import metaflow.datastore.inputs
22
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
 
@@ -728,136 +728,84 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
728
728
  ...
729
729
 
730
730
  @typing.overload
731
- 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]]]:
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]]]:
732
732
  """
733
- Specifies the resources needed when executing this step.
734
-
735
- Use `@resources` to specify the resource requirements
736
- independently of the specific compute layer (`@batch`, `@kubernetes`).
733
+ Creates a human-readable report, a Metaflow Card, after this step completes.
737
734
 
738
- You can choose the compute layer on the command line by executing e.g.
739
- ```
740
- python myflow.py run --with batch
741
- ```
742
- or
743
- ```
744
- python myflow.py run --with kubernetes
745
- ```
746
- which executes the flow on the desired system using the
747
- requirements specified in `@resources`.
735
+ Note that you may add multiple `@card` decorators in a step with different parameters.
748
736
 
749
737
  Parameters
750
738
  ----------
751
- cpu : int, default 1
752
- Number of CPUs required for this step.
753
- gpu : int, default 0
754
- Number of GPUs required for this step.
755
- disk : int, optional, default None
756
- Disk size (in MB) required for this step. Only applies on Kubernetes.
757
- memory : int, default 4096
758
- Memory size (in MB) required for this step.
759
- shared_memory : int, optional, default None
760
- The value for the size (in MiB) of the /dev/shm volume for this step.
761
- This parameter maps to the `--shm-size` option in Docker.
739
+ type : str, default 'default'
740
+ Card type.
741
+ id : str, optional, default None
742
+ If multiple cards are present, use this id to identify this card.
743
+ options : Dict[str, Any], default {}
744
+ Options passed to the card. The contents depend on the card type.
745
+ timeout : int, default 45
746
+ Interrupt reporting if it takes more than this many seconds.
747
+
748
+
762
749
  """
763
750
  ...
764
751
 
765
752
  @typing.overload
766
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
753
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
767
754
  ...
768
755
 
769
756
  @typing.overload
770
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
757
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
771
758
  ...
772
759
 
773
- 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):
760
+ def card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, type: str = "default", id: typing.Optional[str] = None, options: typing.Dict[str, typing.Any] = {}, timeout: int = 45):
774
761
  """
775
- Specifies the resources needed when executing this step.
776
-
777
- Use `@resources` to specify the resource requirements
778
- independently of the specific compute layer (`@batch`, `@kubernetes`).
762
+ Creates a human-readable report, a Metaflow Card, after this step completes.
779
763
 
780
- You can choose the compute layer on the command line by executing e.g.
781
- ```
782
- python myflow.py run --with batch
783
- ```
784
- or
785
- ```
786
- python myflow.py run --with kubernetes
787
- ```
788
- which executes the flow on the desired system using the
789
- requirements specified in `@resources`.
764
+ Note that you may add multiple `@card` decorators in a step with different parameters.
790
765
 
791
766
  Parameters
792
767
  ----------
793
- cpu : int, default 1
794
- Number of CPUs required for this step.
795
- gpu : int, default 0
796
- Number of GPUs required for this step.
797
- disk : int, optional, default None
798
- Disk size (in MB) required for this step. Only applies on Kubernetes.
799
- memory : int, default 4096
800
- Memory size (in MB) required for this step.
801
- shared_memory : int, optional, default None
802
- The value for the size (in MiB) of the /dev/shm volume for this step.
803
- This parameter maps to the `--shm-size` option in Docker.
768
+ type : str, default 'default'
769
+ Card type.
770
+ id : str, optional, default None
771
+ If multiple cards are present, use this id to identify this card.
772
+ options : Dict[str, Any], default {}
773
+ Options passed to the card. The contents depend on the card type.
774
+ timeout : int, default 45
775
+ Interrupt reporting if it takes more than this many seconds.
776
+
777
+
804
778
  """
805
779
  ...
806
780
 
807
781
  @typing.overload
808
- 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]]]:
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]]]:
809
783
  """
810
- Specifies a timeout for your step.
811
-
812
- This decorator is useful if this step may hang indefinitely.
813
-
814
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
815
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
816
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
817
-
818
- Note that all the values specified in parameters are added together so if you specify
819
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
784
+ Specifies environment variables to be set prior to the execution of a step.
820
785
 
821
786
  Parameters
822
787
  ----------
823
- seconds : int, default 0
824
- Number of seconds to wait prior to timing out.
825
- minutes : int, default 0
826
- Number of minutes to wait prior to timing out.
827
- hours : int, default 0
828
- Number of hours to wait prior to timing out.
788
+ vars : Dict[str, str], default {}
789
+ Dictionary of environment variables to set.
829
790
  """
830
791
  ...
831
792
 
832
793
  @typing.overload
833
- def timeout(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]:
834
795
  ...
835
796
 
836
797
  @typing.overload
837
- def timeout(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]:
838
799
  ...
839
800
 
840
- 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):
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] = {}):
841
802
  """
842
- Specifies a timeout for your step.
843
-
844
- This decorator is useful if this step may hang indefinitely.
845
-
846
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
847
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
848
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
849
-
850
- Note that all the values specified in parameters are added together so if you specify
851
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
803
+ Specifies environment variables to be set prior to the execution of a step.
852
804
 
853
805
  Parameters
854
806
  ----------
855
- seconds : int, default 0
856
- Number of seconds to wait prior to timing out.
857
- minutes : int, default 0
858
- Number of minutes to wait prior to timing out.
859
- hours : int, default 0
860
- Number of hours to wait prior to timing out.
807
+ vars : Dict[str, str], default {}
808
+ Dictionary of environment variables to set.
861
809
  """
862
810
  ...
863
811
 
@@ -911,229 +859,206 @@ def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
911
859
  ...
912
860
 
913
861
  @typing.overload
914
- 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]]]:
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]]]:
915
863
  """
916
- Specifies that this step should execute on [AWS Batch](https://aws.amazon.com/batch/).
864
+ Specifies the PyPI packages for the step.
865
+
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.
917
870
 
918
871
  Parameters
919
872
  ----------
920
- cpu : int, default 1
921
- Number of CPUs required for this step. If `@resources` is
922
- also present, the maximum value from all decorators is used.
923
- gpu : int, default 0
924
- Number of GPUs required for this step. If `@resources` is
925
- also present, the maximum value from all decorators is used.
926
- memory : int, default 4096
927
- Memory size (in MB) required for this step. If
928
- `@resources` is also present, the maximum value from all decorators is
929
- used.
930
- image : str, optional, default None
931
- Docker image to use when launching on AWS Batch. If not specified, and
932
- METAFLOW_BATCH_CONTAINER_IMAGE is specified, that image is used. If
933
- not, a default Docker image mapping to the current version of Python is used.
934
- queue : str, default METAFLOW_BATCH_JOB_QUEUE
935
- AWS Batch Job Queue to submit the job to.
936
- iam_role : str, default METAFLOW_ECS_S3_ACCESS_IAM_ROLE
937
- AWS IAM role that AWS Batch container uses to access AWS cloud resources.
938
- execution_role : str, default METAFLOW_ECS_FARGATE_EXECUTION_ROLE
939
- AWS IAM role that AWS Batch can use [to trigger AWS Fargate tasks]
940
- (https://docs.aws.amazon.com/batch/latest/userguide/execution-IAM-role.html).
941
- shared_memory : int, optional, default None
942
- The value for the size (in MiB) of the /dev/shm volume for this step.
943
- This parameter maps to the `--shm-size` option in Docker.
944
- max_swap : int, optional, default None
945
- The total amount of swap memory (in MiB) a container can use for this
946
- step. This parameter is translated to the `--memory-swap` option in
947
- Docker where the value is the sum of the container memory plus the
948
- `max_swap` value.
949
- swappiness : int, optional, default None
950
- This allows you to tune memory swappiness behavior for this step.
951
- A swappiness value of 0 causes swapping not to happen unless absolutely
952
- necessary. A swappiness value of 100 causes pages to be swapped very
953
- aggressively. Accepted values are whole numbers between 0 and 100.
954
- use_tmpfs : bool, default False
955
- This enables an explicit tmpfs mount for this step. Note that tmpfs is
956
- not available on Fargate compute environments
957
- tmpfs_tempdir : bool, default True
958
- sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
959
- tmpfs_size : int, optional, default None
960
- The value for the size (in MiB) of the tmpfs mount for this step.
961
- This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
962
- memory allocated for this step.
963
- tmpfs_path : str, optional, default None
964
- Path to tmpfs mount for this step. Defaults to /metaflow_temp.
965
- inferentia : int, default 0
966
- Number of Inferentia chips required for this step.
967
- trainium : int, default None
968
- Alias for inferentia. Use only one of the two.
969
- efa : int, default 0
970
- Number of elastic fabric adapter network devices to attach to container
971
- ephemeral_storage : int, default None
972
- The total amount, in GiB, of ephemeral storage to set for the task, 21-200GiB.
973
- This is only relevant for Fargate compute environments
974
- log_driver: str, optional, default None
975
- The log driver to use for the Amazon ECS container.
976
- log_options: List[str], optional, default None
977
- List of strings containing options for the chosen log driver. The configurable values
978
- depend on the `log driver` chosen. Validation of these options is not supported yet.
979
- Example: [`awslogs-group:aws/batch/job`]
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.
980
879
  """
981
880
  ...
982
881
 
983
882
  @typing.overload
984
- def batch(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]:
985
884
  ...
986
885
 
987
886
  @typing.overload
988
- def batch(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]:
989
888
  ...
990
889
 
991
- 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):
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):
992
891
  """
993
- Specifies that this step should execute on [AWS Batch](https://aws.amazon.com/batch/).
892
+ Specifies the PyPI packages for the step.
893
+
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.
994
898
 
995
899
  Parameters
996
900
  ----------
997
- cpu : int, default 1
998
- Number of CPUs required for this step. If `@resources` is
999
- also present, the maximum value from all decorators is used.
1000
- gpu : int, default 0
1001
- Number of GPUs required for this step. If `@resources` is
1002
- also present, the maximum value from all decorators is used.
1003
- memory : int, default 4096
1004
- Memory size (in MB) required for this step. If
1005
- `@resources` is also present, the maximum value from all decorators is
1006
- used.
1007
- image : str, optional, default None
1008
- Docker image to use when launching on AWS Batch. If not specified, and
1009
- METAFLOW_BATCH_CONTAINER_IMAGE is specified, that image is used. If
1010
- not, a default Docker image mapping to the current version of Python is used.
1011
- queue : str, default METAFLOW_BATCH_JOB_QUEUE
1012
- AWS Batch Job Queue to submit the job to.
1013
- iam_role : str, default METAFLOW_ECS_S3_ACCESS_IAM_ROLE
1014
- AWS IAM role that AWS Batch container uses to access AWS cloud resources.
1015
- execution_role : str, default METAFLOW_ECS_FARGATE_EXECUTION_ROLE
1016
- AWS IAM role that AWS Batch can use [to trigger AWS Fargate tasks]
1017
- (https://docs.aws.amazon.com/batch/latest/userguide/execution-IAM-role.html).
1018
- shared_memory : int, optional, default None
1019
- The value for the size (in MiB) of the /dev/shm volume for this step.
1020
- This parameter maps to the `--shm-size` option in Docker.
1021
- max_swap : int, optional, default None
1022
- The total amount of swap memory (in MiB) a container can use for this
1023
- step. This parameter is translated to the `--memory-swap` option in
1024
- Docker where the value is the sum of the container memory plus the
1025
- `max_swap` value.
1026
- swappiness : int, optional, default None
1027
- This allows you to tune memory swappiness behavior for this step.
1028
- A swappiness value of 0 causes swapping not to happen unless absolutely
1029
- necessary. A swappiness value of 100 causes pages to be swapped very
1030
- aggressively. Accepted values are whole numbers between 0 and 100.
1031
- use_tmpfs : bool, default False
1032
- This enables an explicit tmpfs mount for this step. Note that tmpfs is
1033
- not available on Fargate compute environments
1034
- tmpfs_tempdir : bool, default True
1035
- sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
1036
- tmpfs_size : int, optional, default None
1037
- The value for the size (in MiB) of the tmpfs mount for this step.
1038
- This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
1039
- memory allocated for this step.
1040
- tmpfs_path : str, optional, default None
1041
- Path to tmpfs mount for this step. Defaults to /metaflow_temp.
1042
- inferentia : int, default 0
1043
- Number of Inferentia chips required for this step.
1044
- trainium : int, default None
1045
- Alias for inferentia. Use only one of the two.
1046
- efa : int, default 0
1047
- Number of elastic fabric adapter network devices to attach to container
1048
- ephemeral_storage : int, default None
1049
- The total amount, in GiB, of ephemeral storage to set for the task, 21-200GiB.
1050
- This is only relevant for Fargate compute environments
1051
- log_driver: str, optional, default None
1052
- The log driver to use for the Amazon ECS container.
1053
- log_options: List[str], optional, default None
1054
- List of strings containing options for the chosen log driver. The configurable values
1055
- depend on the `log driver` chosen. Validation of these options is not supported yet.
1056
- Example: [`awslogs-group:aws/batch/job`]
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.
1057
907
  """
1058
908
  ...
1059
909
 
1060
910
  @typing.overload
1061
- 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]]]:
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]]]:
1062
912
  """
1063
- Specifies environment variables to be set prior to 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`.
1064
928
 
1065
929
  Parameters
1066
930
  ----------
1067
- vars : Dict[str, str], default {}
1068
- Dictionary of environment variables to set.
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.
1069
942
  """
1070
943
  ...
1071
944
 
1072
945
  @typing.overload
1073
- def environment(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]:
1074
947
  ...
1075
948
 
1076
949
  @typing.overload
1077
- def environment(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]:
1078
951
  ...
1079
952
 
1080
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
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):
1081
954
  """
1082
- Specifies environment variables to be set prior to 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`.
1083
970
 
1084
971
  Parameters
1085
972
  ----------
1086
- vars : Dict[str, str], default {}
1087
- Dictionary of environment variables to set.
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.
1088
984
  """
1089
985
  ...
1090
986
 
1091
987
  @typing.overload
1092
- 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]]]:
988
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1093
989
  """
1094
- Specifies the PyPI packages for the step.
1095
-
1096
- Information in this decorator will augment any
1097
- attributes set in the `@pyi_base` flow-level decorator. Hence,
1098
- you can use `@pypi_base` to set packages required by all
1099
- steps and use `@pypi` to specify step-specific overrides.
1100
-
1101
- Parameters
1102
- ----------
1103
- packages : Dict[str, str], default: {}
1104
- Packages to use for this step. The key is the name of the package
1105
- and the value is the version to use.
1106
- python : str, optional, default: None
1107
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1108
- that the version used will correspond to the version of the Python interpreter used to start the run.
990
+ Decorator prototype for all step decorators. This function gets specialized
991
+ and imported for all decorators types by _import_plugin_decorators().
1109
992
  """
1110
993
  ...
1111
994
 
1112
995
  @typing.overload
1113
- def pypi(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]:
1114
997
  ...
1115
998
 
1116
- @typing.overload
1117
- def pypi(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
+ """
1118
1004
  ...
1119
1005
 
1120
- 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):
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]]]:
1121
1007
  """
1122
- Specifies the PyPI packages for the step.
1123
-
1124
- Information in this decorator will augment any
1125
- attributes set in the `@pyi_base` flow-level decorator. Hence,
1126
- you can use `@pypi_base` to set packages required by all
1127
- steps and use `@pypi` to specify step-specific overrides.
1008
+ Specifies that this step should execute on Kubernetes.
1128
1009
 
1129
1010
  Parameters
1130
1011
  ----------
1131
- packages : Dict[str, str], default: {}
1132
- Packages to use for this step. The key is the name of the package
1133
- and the value is the version to use.
1134
- python : str, optional, default: None
1135
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1136
- that the version used will correspond to the version of the Python interpreter used to start the run.
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
1137
1062
  """
1138
1063
  ...
1139
1064
 
@@ -1195,53 +1120,206 @@ def conda(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
1195
1120
  ...
1196
1121
 
1197
1122
  @typing.overload
1198
- 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]]]:
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]]]:
1199
1124
  """
1200
- Creates a human-readable report, a Metaflow Card, after this step completes.
1125
+ Specifies a timeout for your step.
1201
1126
 
1202
- Note that you may add multiple `@card` decorators in a step with different parameters.
1127
+ This decorator is useful if this step may hang indefinitely.
1203
1128
 
1204
- Parameters
1205
- ----------
1206
- type : str, default 'default'
1207
- Card type.
1208
- id : str, optional, default None
1209
- If multiple cards are present, use this id to identify this card.
1210
- options : Dict[str, Any], default {}
1211
- Options passed to the card. The contents depend on the card type.
1212
- timeout : int, default 45
1213
- Interrupt reporting if it takes more than this many seconds.
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.
1214
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.
1215
1135
 
1136
+ Parameters
1137
+ ----------
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.
1216
1144
  """
1217
1145
  ...
1218
1146
 
1219
1147
  @typing.overload
1220
- def card(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]:
1221
1149
  ...
1222
1150
 
1223
1151
  @typing.overload
1224
- def card(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]:
1225
1153
  ...
1226
1154
 
1227
- def card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, type: str = "default", id: typing.Optional[str] = None, options: typing.Dict[str, typing.Any] = {}, timeout: int = 45):
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):
1228
1156
  """
1229
- Creates a human-readable report, a Metaflow Card, after this step completes.
1157
+ Specifies a timeout for your step.
1230
1158
 
1231
- Note that you may add multiple `@card` decorators in a step with different parameters.
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.
1232
1167
 
1233
1168
  Parameters
1234
1169
  ----------
1235
- type : str, default 'default'
1236
- Card type.
1237
- id : str, optional, default None
1238
- If multiple cards are present, use this id to identify this card.
1239
- options : Dict[str, Any], default {}
1240
- Options passed to the card. The contents depend on the card type.
1241
- timeout : int, default 45
1242
- Interrupt reporting if it takes more than this many seconds.
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.
1176
+ """
1177
+ ...
1178
+
1179
+ @typing.overload
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]]]:
1181
+ """
1182
+ Specifies that this step should execute on [AWS Batch](https://aws.amazon.com/batch/).
1243
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`]
1246
+ """
1247
+ ...
1248
+
1249
+ @typing.overload
1250
+ def batch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1251
+ ...
1252
+
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):
1258
+ """
1259
+ Specifies that this step should execute on [AWS Batch](https://aws.amazon.com/batch/).
1244
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`]
1245
1323
  """
1246
1324
  ...
1247
1325
 
@@ -1278,65 +1356,6 @@ def secrets(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None],
1278
1356
  """
1279
1357
  ...
1280
1358
 
1281
- 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]]]:
1282
- """
1283
- Specifies that this step should execute on Kubernetes.
1284
-
1285
- Parameters
1286
- ----------
1287
- cpu : int, default 1
1288
- Number of CPUs required for this step. If `@resources` is
1289
- also present, the maximum value from all decorators is used.
1290
- memory : int, default 4096
1291
- Memory size (in MB) required for this step. If
1292
- `@resources` is also present, the maximum value from all decorators is
1293
- used.
1294
- disk : int, default 10240
1295
- Disk size (in MB) required for this step. If
1296
- `@resources` is also present, the maximum value from all decorators is
1297
- used.
1298
- image : str, optional, default None
1299
- Docker image to use when launching on Kubernetes. If not specified, and
1300
- METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
1301
- not, a default Docker image mapping to the current version of Python is used.
1302
- image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
1303
- If given, the imagePullPolicy to be applied to the Docker image of the step.
1304
- service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
1305
- Kubernetes service account to use when launching pod in Kubernetes.
1306
- secrets : List[str], optional, default None
1307
- Kubernetes secrets to use when launching pod in Kubernetes. These
1308
- secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
1309
- in Metaflow configuration.
1310
- namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
1311
- Kubernetes namespace to use when launching pod in Kubernetes.
1312
- gpu : int, optional, default None
1313
- Number of GPUs required for this step. A value of zero implies that
1314
- the scheduled node should not have GPUs.
1315
- gpu_vendor : str, default KUBERNETES_GPU_VENDOR
1316
- The vendor of the GPUs to be used for this step.
1317
- tolerations : List[str], default []
1318
- The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
1319
- Kubernetes tolerations to use when launching pod in Kubernetes.
1320
- use_tmpfs : bool, default False
1321
- This enables an explicit tmpfs mount for this step.
1322
- tmpfs_tempdir : bool, default True
1323
- sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
1324
- tmpfs_size : int, optional, default: None
1325
- The value for the size (in MiB) of the tmpfs mount for this step.
1326
- This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
1327
- memory allocated for this step.
1328
- tmpfs_path : str, optional, default /metaflow_temp
1329
- Path to tmpfs mount for this step.
1330
- persistent_volume_claims : Dict[str, str], optional, default None
1331
- A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
1332
- volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
1333
- shared_memory: int, optional
1334
- Shared memory size (in MiB) required for this step
1335
- port: int, optional
1336
- Port number to specify in the Kubernetes job object
1337
- """
1338
- ...
1339
-
1340
1359
  @typing.overload
1341
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]]]:
1342
1361
  """
@@ -1439,104 +1458,19 @@ def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly:
1439
1458
  """
1440
1459
  ...
1441
1460
 
1442
- @typing.overload
1443
- 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]]:
1444
- """
1445
- Specifies the flow(s) that this flow depends on.
1446
-
1447
- ```
1448
- @trigger_on_finish(flow='FooFlow')
1449
- ```
1450
- or
1451
- ```
1452
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1453
- ```
1454
- This decorator respects the @project decorator and triggers the flow
1455
- when upstream runs within the same namespace complete successfully
1456
-
1457
- Additionally, you can specify project aware upstream flow dependencies
1458
- by specifying the fully qualified project_flow_name.
1459
- ```
1460
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1461
- ```
1462
- or
1463
- ```
1464
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1465
- ```
1466
-
1467
- You can also specify just the project or project branch (other values will be
1468
- inferred from the current project or project branch):
1469
- ```
1470
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1471
- ```
1472
-
1473
- Note that `branch` is typically one of:
1474
- - `prod`
1475
- - `user.bob`
1476
- - `test.my_experiment`
1477
- - `prod.staging`
1478
-
1479
- Parameters
1480
- ----------
1481
- flow : Union[str, Dict[str, str]], optional, default None
1482
- Upstream flow dependency for this flow.
1483
- flows : List[Union[str, Dict[str, str]]], default []
1484
- Upstream flow dependencies for this flow.
1485
- options : Dict[str, Any], default {}
1486
- Backend-specific configuration for tuning eventing behavior.
1487
-
1488
-
1489
- """
1490
- ...
1491
-
1492
- @typing.overload
1493
- def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1494
- ...
1495
-
1496
- 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] = {}):
1461
+ def project(*, name: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1497
1462
  """
1498
- Specifies the flow(s) that this flow depends on.
1499
-
1500
- ```
1501
- @trigger_on_finish(flow='FooFlow')
1502
- ```
1503
- or
1504
- ```
1505
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1506
- ```
1507
- This decorator respects the @project decorator and triggers the flow
1508
- when upstream runs within the same namespace complete successfully
1509
-
1510
- Additionally, you can specify project aware upstream flow dependencies
1511
- by specifying the fully qualified project_flow_name.
1512
- ```
1513
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1514
- ```
1515
- or
1516
- ```
1517
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1518
- ```
1519
-
1520
- You can also specify just the project or project branch (other values will be
1521
- inferred from the current project or project branch):
1522
- ```
1523
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1524
- ```
1463
+ Specifies what flows belong to the same project.
1525
1464
 
1526
- Note that `branch` is typically one of:
1527
- - `prod`
1528
- - `user.bob`
1529
- - `test.my_experiment`
1530
- - `prod.staging`
1465
+ A project-specific namespace is created for all flows that
1466
+ use the same `@project(name)`.
1531
1467
 
1532
1468
  Parameters
1533
1469
  ----------
1534
- flow : Union[str, Dict[str, str]], optional, default None
1535
- Upstream flow dependency for this flow.
1536
- flows : List[Union[str, Dict[str, str]]], default []
1537
- Upstream flow dependencies for this flow.
1538
- options : Dict[str, Any], default {}
1539
- 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.
1540
1474
 
1541
1475
 
1542
1476
  """
@@ -1637,45 +1571,42 @@ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: t
1637
1571
  """
1638
1572
  ...
1639
1573
 
1640
- 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]]:
1574
+ @typing.overload
1575
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1641
1576
  """
1642
- 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.
1643
- 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.
1644
-
1645
- Parameters
1646
- ----------
1647
- timeout : int
1648
- Time, in seconds before the task times out and fails. (Default: 3600)
1649
- poke_interval : int
1650
- Time in seconds that the job should wait in between each try. (Default: 60)
1651
- mode : str
1652
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1653
- exponential_backoff : bool
1654
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1655
- pool : str
1656
- the slot pool this task should run in,
1657
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1658
- soft_fail : bool
1659
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1660
- name : str
1661
- Name of the sensor on Airflow
1662
- description : str
1663
- Description of sensor in the Airflow UI
1664
- external_dag_id : str
1665
- The dag_id that contains the task you want to wait for.
1666
- external_task_ids : List[str]
1667
- The list of task_ids that you want to wait for.
1668
- If None (default value) the sensor waits for the DAG. (Default: None)
1669
- allowed_states : List[str]
1670
- Iterable of allowed states, (Default: ['success'])
1671
- failed_states : List[str]
1672
- Iterable of failed or dis-allowed states. (Default: None)
1673
- execution_delta : datetime.timedelta
1674
- time difference with the previous execution to look at,
1675
- the default is the same logical date as the current task or DAG. (Default: None)
1676
- check_existence: bool
1677
- Set to True to check if the external task exists or check if
1678
- the DAG to wait for exists. (Default: True)
1577
+ Specifies the PyPI packages for all steps of the flow.
1578
+
1579
+ Use `@pypi_base` to set common packages required by all
1580
+ steps and use `@pypi` to specify step-specific overrides.
1581
+ Parameters
1582
+ ----------
1583
+ packages : Dict[str, str], default: {}
1584
+ Packages to use for this flow. The key is the name of the package
1585
+ and the value is the version to use.
1586
+ python : str, optional, default: None
1587
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1588
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1589
+ """
1590
+ ...
1591
+
1592
+ @typing.overload
1593
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1594
+ ...
1595
+
1596
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1597
+ """
1598
+ Specifies the PyPI packages for all steps of the flow.
1599
+
1600
+ Use `@pypi_base` to set common packages required by all
1601
+ steps and use `@pypi` to specify step-specific overrides.
1602
+ Parameters
1603
+ ----------
1604
+ packages : Dict[str, str], default: {}
1605
+ Packages to use for this flow. The key is the name of the package
1606
+ and the value is the version to use.
1607
+ python : str, optional, default: None
1608
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1609
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1679
1610
  """
1680
1611
  ...
1681
1612
 
@@ -1728,63 +1659,6 @@ def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packa
1728
1659
  """
1729
1660
  ...
1730
1661
 
1731
- def project(*, name: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1732
- """
1733
- Specifies what flows belong to the same project.
1734
-
1735
- A project-specific namespace is created for all flows that
1736
- use the same `@project(name)`.
1737
-
1738
- Parameters
1739
- ----------
1740
- name : str
1741
- Project name. Make sure that the name is unique amongst all
1742
- projects that use the same production scheduler. The name may
1743
- contain only lowercase alphanumeric characters and underscores.
1744
-
1745
-
1746
- """
1747
- ...
1748
-
1749
- @typing.overload
1750
- def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1751
- """
1752
- Specifies the PyPI packages for all steps of the flow.
1753
-
1754
- Use `@pypi_base` to set common packages required by all
1755
- steps and use `@pypi` to specify step-specific overrides.
1756
- Parameters
1757
- ----------
1758
- packages : Dict[str, str], default: {}
1759
- Packages to use for this flow. The key is the name of the package
1760
- and the value is the version to use.
1761
- python : str, optional, default: None
1762
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1763
- that the version used will correspond to the version of the Python interpreter used to start the run.
1764
- """
1765
- ...
1766
-
1767
- @typing.overload
1768
- def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1769
- ...
1770
-
1771
- def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1772
- """
1773
- Specifies the PyPI packages for all steps of the flow.
1774
-
1775
- Use `@pypi_base` to set common packages required by all
1776
- steps and use `@pypi` to specify step-specific overrides.
1777
- Parameters
1778
- ----------
1779
- packages : Dict[str, str], default: {}
1780
- Packages to use for this flow. The key is the name of the package
1781
- and the value is the version to use.
1782
- python : str, optional, default: None
1783
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1784
- that the version used will correspond to the version of the Python interpreter used to start the run.
1785
- """
1786
- ...
1787
-
1788
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]]:
1789
1663
  """
1790
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)
@@ -1827,6 +1701,151 @@ def airflow_s3_key_sensor(*, timeout: int, poke_interval: int, mode: str, expone
1827
1701
  """
1828
1702
  ...
1829
1703
 
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]]:
1705
+ """
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.
1707
+ 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.
1708
+
1709
+ Parameters
1710
+ ----------
1711
+ timeout : int
1712
+ Time, in seconds before the task times out and fails. (Default: 3600)
1713
+ poke_interval : int
1714
+ Time in seconds that the job should wait in between each try. (Default: 60)
1715
+ mode : str
1716
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1717
+ exponential_backoff : bool
1718
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1719
+ pool : str
1720
+ the slot pool this task should run in,
1721
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1722
+ soft_fail : bool
1723
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1724
+ name : str
1725
+ Name of the sensor on Airflow
1726
+ description : str
1727
+ Description of sensor in the Airflow UI
1728
+ external_dag_id : str
1729
+ The dag_id that contains the task you want to wait for.
1730
+ external_task_ids : List[str]
1731
+ The list of task_ids that you want to wait for.
1732
+ If None (default value) the sensor waits for the DAG. (Default: None)
1733
+ allowed_states : List[str]
1734
+ Iterable of allowed states, (Default: ['success'])
1735
+ failed_states : List[str]
1736
+ Iterable of failed or dis-allowed states. (Default: None)
1737
+ execution_delta : datetime.timedelta
1738
+ time difference with the previous execution to look at,
1739
+ the default is the same logical date as the current task or DAG. (Default: None)
1740
+ check_existence: bool
1741
+ Set to True to check if the external task exists or check if
1742
+ the DAG to wait for exists. (Default: True)
1743
+ """
1744
+ ...
1745
+
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]]:
1748
+ """
1749
+ Specifies the flow(s) that this flow depends on.
1750
+
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
1760
+
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
+ ```
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
+ ```
1776
+
1777
+ Note that `branch` is typically one of:
1778
+ - `prod`
1779
+ - `user.bob`
1780
+ - `test.my_experiment`
1781
+ - `prod.staging`
1782
+
1783
+ Parameters
1784
+ ----------
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
+
1793
+ """
1794
+ ...
1795
+
1796
+ @typing.overload
1797
+ def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1798
+ ...
1799
+
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] = {}):
1801
+ """
1802
+ Specifies the flow(s) that this flow depends on.
1803
+
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`
1835
+
1836
+ Parameters
1837
+ ----------
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
+ """
1847
+ ...
1848
+
1830
1849
  def namespace(ns: typing.Optional[str]) -> typing.Optional[str]:
1831
1850
  """
1832
1851
  Switch namespace to the one provided.
@@ -2975,3 +2994,33 @@ class NBRunner(object, metaclass=type):
2975
2994
  ...
2976
2995
  ...
2977
2996
 
2997
+ class Deployer(object, metaclass=type):
2998
+ def __init__(self, flow_file: str, show_output: bool = True, profile: typing.Optional[str] = None, env: typing.Optional[typing.Dict] = None, cwd: typing.Optional[str] = None, **kwargs):
2999
+ ...
3000
+ def _Deployer__make_function(self, deployer_class):
3001
+ """
3002
+ Create a function for the given deployer class.
3003
+
3004
+ Parameters
3005
+ ----------
3006
+ deployer_class : Type[DeployerImpl]
3007
+ Deployer implementation class.
3008
+
3009
+ Returns
3010
+ -------
3011
+ Callable
3012
+ Function that initializes and returns an instance of the deployer class.
3013
+ """
3014
+ ...
3015
+ ...
3016
+
3017
+ class NBDeployer(object, metaclass=type):
3018
+ def __init__(self, flow, show_output: bool = True, profile: typing.Optional[str] = None, env: typing.Optional[typing.Dict] = None, base_dir: str = "/tmp", **kwargs):
3019
+ ...
3020
+ def cleanup(self):
3021
+ """
3022
+ Delete any temporary files created during execution.
3023
+ """
3024
+ ...
3025
+ ...
3026
+