ob-metaflow-stubs 6.0.3.108__py2.py3-none-any.whl → 6.0.3.109rc0__py2.py3-none-any.whl

Sign up to get free protection for your applications and to get access to all the features.
Files changed (140) hide show
  1. metaflow-stubs/__init__.pyi +501 -487
  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 +7 -7
  6. metaflow-stubs/client/filecache.pyi +3 -3
  7. metaflow-stubs/clone_util.pyi +2 -2
  8. metaflow-stubs/events.pyi +3 -3
  9. metaflow-stubs/exception.pyi +2 -2
  10. metaflow-stubs/flowspec.pyi +5 -5
  11. metaflow-stubs/generated_for.txt +1 -1
  12. metaflow-stubs/includefile.pyi +4 -4
  13. metaflow-stubs/info_file.pyi +2 -2
  14. metaflow-stubs/metadata/metadata.pyi +2 -2
  15. metaflow-stubs/metadata/util.pyi +2 -2
  16. metaflow-stubs/metaflow_config.pyi +2 -2
  17. metaflow-stubs/metaflow_current.pyi +6 -6
  18. metaflow-stubs/mflog/mflog.pyi +2 -2
  19. metaflow-stubs/multicore_utils.pyi +2 -2
  20. metaflow-stubs/parameters.pyi +4 -4
  21. metaflow-stubs/plugins/__init__.pyi +2 -2
  22. metaflow-stubs/plugins/airflow/__init__.pyi +2 -2
  23. metaflow-stubs/plugins/airflow/airflow_utils.pyi +2 -2
  24. metaflow-stubs/plugins/airflow/exception.pyi +2 -2
  25. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +2 -2
  26. metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi +3 -3
  27. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +3 -3
  28. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +3 -3
  29. metaflow-stubs/plugins/argo/__init__.pyi +2 -2
  30. metaflow-stubs/plugins/argo/argo_client.pyi +2 -2
  31. metaflow-stubs/plugins/argo/argo_events.pyi +2 -2
  32. metaflow-stubs/plugins/argo/argo_workflows.pyi +5 -5
  33. metaflow-stubs/plugins/argo/argo_workflows_cli.pyi +7 -7
  34. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +7 -7
  35. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +2 -2
  36. metaflow-stubs/plugins/aws/__init__.pyi +2 -2
  37. metaflow-stubs/plugins/aws/aws_client.pyi +2 -2
  38. metaflow-stubs/plugins/aws/aws_utils.pyi +2 -2
  39. metaflow-stubs/plugins/aws/batch/__init__.pyi +2 -2
  40. metaflow-stubs/plugins/aws/batch/batch.pyi +2 -2
  41. metaflow-stubs/plugins/aws/batch/batch_client.pyi +2 -2
  42. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +3 -3
  43. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +2 -2
  44. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +3 -3
  45. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +2 -2
  46. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +2 -2
  47. metaflow-stubs/plugins/aws/step_functions/production_token.pyi +2 -2
  48. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +2 -2
  49. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +2 -2
  50. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +2 -2
  51. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer.pyi +2 -2
  52. metaflow-stubs/plugins/azure/__init__.pyi +2 -2
  53. metaflow-stubs/plugins/azure/azure_credential.pyi +2 -2
  54. metaflow-stubs/plugins/azure/azure_exceptions.pyi +2 -2
  55. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +3 -3
  56. metaflow-stubs/plugins/azure/azure_utils.pyi +2 -2
  57. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +2 -2
  58. metaflow-stubs/plugins/azure/includefile_support.pyi +2 -2
  59. metaflow-stubs/plugins/cards/__init__.pyi +2 -2
  60. metaflow-stubs/plugins/cards/card_cli.pyi +4 -4
  61. metaflow-stubs/plugins/cards/card_client.pyi +3 -3
  62. metaflow-stubs/plugins/cards/card_creator.pyi +2 -2
  63. metaflow-stubs/plugins/cards/card_datastore.pyi +2 -2
  64. metaflow-stubs/plugins/cards/card_decorator.pyi +2 -2
  65. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +2 -2
  66. metaflow-stubs/plugins/cards/card_modules/basic.pyi +4 -4
  67. metaflow-stubs/plugins/cards/card_modules/card.pyi +2 -2
  68. metaflow-stubs/plugins/cards/card_modules/components.pyi +4 -4
  69. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +2 -2
  70. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +2 -2
  71. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +3 -3
  72. metaflow-stubs/plugins/cards/card_resolver.pyi +2 -2
  73. metaflow-stubs/plugins/cards/component_serializer.pyi +4 -4
  74. metaflow-stubs/plugins/cards/exception.pyi +2 -2
  75. metaflow-stubs/plugins/catch_decorator.pyi +3 -3
  76. metaflow-stubs/plugins/datatools/__init__.pyi +3 -3
  77. metaflow-stubs/plugins/datatools/local.pyi +2 -2
  78. metaflow-stubs/plugins/datatools/s3/__init__.pyi +3 -3
  79. metaflow-stubs/plugins/datatools/s3/s3.pyi +5 -5
  80. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +2 -2
  81. metaflow-stubs/plugins/datatools/s3/s3util.pyi +2 -2
  82. metaflow-stubs/plugins/debug_logger.pyi +2 -2
  83. metaflow-stubs/plugins/debug_monitor.pyi +2 -2
  84. metaflow-stubs/plugins/environment_decorator.pyi +2 -2
  85. metaflow-stubs/plugins/events_decorator.pyi +2 -2
  86. metaflow-stubs/plugins/frameworks/__init__.pyi +2 -2
  87. metaflow-stubs/plugins/frameworks/pytorch.pyi +2 -2
  88. metaflow-stubs/plugins/gcp/__init__.pyi +2 -2
  89. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +3 -3
  90. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +2 -2
  91. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +2 -2
  92. metaflow-stubs/plugins/gcp/gs_utils.pyi +2 -2
  93. metaflow-stubs/plugins/gcp/includefile_support.pyi +2 -2
  94. metaflow-stubs/plugins/kubernetes/__init__.pyi +2 -2
  95. metaflow-stubs/plugins/kubernetes/kube_utils.pyi +2 -2
  96. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +2 -2
  97. metaflow-stubs/plugins/kubernetes/kubernetes_cli.pyi +2 -2
  98. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +2 -2
  99. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +3 -3
  100. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +2 -2
  101. metaflow-stubs/plugins/logs_cli.pyi +4 -4
  102. metaflow-stubs/plugins/package_cli.pyi +2 -2
  103. metaflow-stubs/plugins/parallel_decorator.pyi +2 -2
  104. metaflow-stubs/plugins/perimeters.pyi +2 -2
  105. metaflow-stubs/plugins/project_decorator.pyi +2 -2
  106. metaflow-stubs/plugins/pypi/__init__.pyi +2 -2
  107. metaflow-stubs/plugins/pypi/conda_decorator.pyi +2 -2
  108. metaflow-stubs/plugins/pypi/conda_environment.pyi +5 -5
  109. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +2 -2
  110. metaflow-stubs/plugins/pypi/pypi_environment.pyi +3 -3
  111. metaflow-stubs/plugins/pypi/utils.pyi +2 -2
  112. metaflow-stubs/plugins/resources_decorator.pyi +2 -2
  113. metaflow-stubs/plugins/retry_decorator.pyi +2 -2
  114. metaflow-stubs/plugins/secrets/__init__.pyi +2 -2
  115. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +2 -2
  116. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +2 -2
  117. metaflow-stubs/plugins/storage_executor.pyi +2 -2
  118. metaflow-stubs/plugins/tag_cli.pyi +6 -6
  119. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +2 -2
  120. metaflow-stubs/plugins/timeout_decorator.pyi +3 -3
  121. metaflow-stubs/procpoll.pyi +2 -2
  122. metaflow-stubs/profilers/__init__.pyi +2 -2
  123. metaflow-stubs/pylint_wrapper.pyi +2 -2
  124. metaflow-stubs/runner/__init__.pyi +2 -2
  125. metaflow-stubs/runner/deployer.pyi +3 -3
  126. metaflow-stubs/runner/metaflow_runner.pyi +6 -6
  127. metaflow-stubs/runner/nbdeploy.pyi +2 -2
  128. metaflow-stubs/runner/nbrun.pyi +2 -2
  129. metaflow-stubs/runner/subprocess_manager.pyi +2 -2
  130. metaflow-stubs/runner/utils.pyi +2 -2
  131. metaflow-stubs/system/__init__.pyi +3 -3
  132. metaflow-stubs/system/system_logger.pyi +2 -2
  133. metaflow-stubs/system/system_monitor.pyi +3 -3
  134. metaflow-stubs/tagging_util.pyi +2 -2
  135. metaflow-stubs/tuple_util.pyi +2 -2
  136. {ob_metaflow_stubs-6.0.3.108.dist-info → ob_metaflow_stubs-6.0.3.109rc0.dist-info}/METADATA +1 -1
  137. ob_metaflow_stubs-6.0.3.109rc0.dist-info/RECORD +140 -0
  138. ob_metaflow_stubs-6.0.3.108.dist-info/RECORD +0 -140
  139. {ob_metaflow_stubs-6.0.3.108.dist-info → ob_metaflow_stubs-6.0.3.109rc0.dist-info}/WHEEL +0 -0
  140. {ob_metaflow_stubs-6.0.3.108.dist-info → ob_metaflow_stubs-6.0.3.109rc0.dist-info}/top_level.txt +0 -0
@@ -1,24 +1,24 @@
1
1
  ##################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
- # MF version: 2.12.25.1+obcheckpoint(0.0.11);ob(v1) #
4
- # Generated on 2024-10-08T12:29:45.085886 #
3
+ # MF version: 2.12.25.1+obcheckpoint(0.0.12);ob(v1) #
4
+ # Generated on 2024-10-10T01:08:08.920544 #
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._vendor.click.types
12
- import metaflow.metaflow_current
13
- import metaflow.datastore.inputs
11
+ import metaflow.events
14
12
  import metaflow_extensions.obcheckpoint.plugins.machine_learning_utilities.datastructures
15
- import typing
16
- import metaflow.flowspec
17
- import datetime
18
13
  import metaflow.runner.metaflow_runner
14
+ import datetime
15
+ import metaflow.metaflow_current
16
+ import typing
17
+ import metaflow._vendor.click.types
19
18
  import metaflow.parameters
20
- import metaflow.events
21
19
  import metaflow.client.core
20
+ import metaflow.datastore.inputs
21
+ import metaflow.flowspec
22
22
  FlowSpecDerived = typing.TypeVar("FlowSpecDerived", bound="FlowSpec", contravariant=False, covariant=False)
23
23
  StepFlag = typing.NewType("StepFlag", bool)
24
24
 
@@ -513,125 +513,19 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
513
513
  ...
514
514
 
515
515
  @typing.overload
516
- 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]]]:
516
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
517
517
  """
518
- Specifies a timeout for your step.
519
-
520
- This decorator is useful if this step may hang indefinitely.
521
-
522
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
523
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
524
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
525
-
526
- Note that all the values specified in parameters are added together so if you specify
527
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
528
-
529
- Parameters
530
- ----------
531
- seconds : int, default 0
532
- Number of seconds to wait prior to timing out.
533
- minutes : int, default 0
534
- Number of minutes to wait prior to timing out.
535
- hours : int, default 0
536
- Number of hours to wait prior to timing out.
518
+ Internal decorator to support Fast bakery
537
519
  """
538
520
  ...
539
521
 
540
522
  @typing.overload
541
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
542
- ...
543
-
544
- @typing.overload
545
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
546
- ...
547
-
548
- 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):
549
- """
550
- Specifies a timeout for your step.
551
-
552
- This decorator is useful if this step may hang indefinitely.
553
-
554
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
555
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
556
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
557
-
558
- Note that all the values specified in parameters are added together so if you specify
559
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
560
-
561
- Parameters
562
- ----------
563
- seconds : int, default 0
564
- Number of seconds to wait prior to timing out.
565
- minutes : int, default 0
566
- Number of minutes to wait prior to timing out.
567
- hours : int, default 0
568
- Number of hours to wait prior to timing out.
569
- """
523
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
570
524
  ...
571
525
 
572
- 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, node_selector: typing.Union[typing.Dict[str, str], str, None] = None, namespace: str = "METAFLOW_KUBERNETES_NAMESPACE", gpu: typing.Optional[int] = None, gpu_vendor: str = "KUBERNETES_GPU_VENDOR", tolerations: typing.List[str] = [], use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = "/metaflow_temp", persistent_volume_claims: typing.Optional[typing.Dict[str, str]] = None, shared_memory: typing.Optional[int] = None, port: typing.Optional[int] = None, compute_pool: typing.Optional[str] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
526
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
573
527
  """
574
- Specifies that this step should execute on Kubernetes.
575
-
576
- Parameters
577
- ----------
578
- cpu : int, default 1
579
- Number of CPUs required for this step. If `@resources` is
580
- also present, the maximum value from all decorators is used.
581
- memory : int, default 4096
582
- Memory size (in MB) required for this step. If
583
- `@resources` is also present, the maximum value from all decorators is
584
- used.
585
- disk : int, default 10240
586
- Disk size (in MB) required for this step. If
587
- `@resources` is also present, the maximum value from all decorators is
588
- used.
589
- image : str, optional, default None
590
- Docker image to use when launching on Kubernetes. If not specified, and
591
- METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
592
- not, a default Docker image mapping to the current version of Python is used.
593
- image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
594
- If given, the imagePullPolicy to be applied to the Docker image of the step.
595
- service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
596
- Kubernetes service account to use when launching pod in Kubernetes.
597
- secrets : List[str], optional, default None
598
- Kubernetes secrets to use when launching pod in Kubernetes. These
599
- secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
600
- in Metaflow configuration.
601
- node_selector: Union[Dict[str,str], str], optional, default None
602
- Kubernetes node selector(s) to apply to the pod running the task.
603
- Can be passed in as a comma separated string of values e.g. "kubernetes.io/os=linux,kubernetes.io/arch=amd64"
604
- or as a dictionary {"kubernetes.io/os": "linux", "kubernetes.io/arch": "amd64"}
605
- namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
606
- Kubernetes namespace to use when launching pod in Kubernetes.
607
- gpu : int, optional, default None
608
- Number of GPUs required for this step. A value of zero implies that
609
- the scheduled node should not have GPUs.
610
- gpu_vendor : str, default KUBERNETES_GPU_VENDOR
611
- The vendor of the GPUs to be used for this step.
612
- tolerations : List[str], default []
613
- The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
614
- Kubernetes tolerations to use when launching pod in Kubernetes.
615
- use_tmpfs : bool, default False
616
- This enables an explicit tmpfs mount for this step.
617
- tmpfs_tempdir : bool, default True
618
- sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
619
- tmpfs_size : int, optional, default: None
620
- The value for the size (in MiB) of the tmpfs mount for this step.
621
- This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
622
- memory allocated for this step.
623
- tmpfs_path : str, optional, default /metaflow_temp
624
- Path to tmpfs mount for this step.
625
- persistent_volume_claims : Dict[str, str], optional, default None
626
- A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
627
- volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
628
- shared_memory: int, optional
629
- Shared memory size (in MiB) required for this step
630
- port: int, optional
631
- Port number to specify in the Kubernetes job object
632
- compute_pool : str, optional, default None
633
- Compute pool to be used for for this step.
634
- If not specified, any accessible compute pool within the perimeter is used.
528
+ Internal decorator to support Fast bakery
635
529
  """
636
530
  ...
637
531
 
@@ -685,51 +579,120 @@ def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
685
579
  ...
686
580
 
687
581
  @typing.overload
688
- 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]]]:
582
+ 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]]]:
689
583
  """
690
- Specifies the PyPI packages for the step.
584
+ Specifies a timeout for your step.
691
585
 
692
- Information in this decorator will augment any
693
- attributes set in the `@pyi_base` flow-level decorator. Hence,
694
- you can use `@pypi_base` to set packages required by all
695
- steps and use `@pypi` to specify step-specific overrides.
586
+ This decorator is useful if this step may hang indefinitely.
587
+
588
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
589
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
590
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
591
+
592
+ Note that all the values specified in parameters are added together so if you specify
593
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
696
594
 
697
595
  Parameters
698
596
  ----------
699
- packages : Dict[str, str], default: {}
700
- Packages to use for this step. The key is the name of the package
701
- and the value is the version to use.
702
- python : str, optional, default: None
703
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
704
- that the version used will correspond to the version of the Python interpreter used to start the run.
597
+ seconds : int, default 0
598
+ Number of seconds to wait prior to timing out.
599
+ minutes : int, default 0
600
+ Number of minutes to wait prior to timing out.
601
+ hours : int, default 0
602
+ Number of hours to wait prior to timing out.
705
603
  """
706
604
  ...
707
605
 
708
606
  @typing.overload
709
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
607
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
710
608
  ...
711
609
 
712
610
  @typing.overload
713
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
611
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
714
612
  ...
715
613
 
716
- 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):
614
+ 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):
717
615
  """
718
- Specifies the PyPI packages for the step.
616
+ Specifies a timeout for your step.
719
617
 
720
- Information in this decorator will augment any
721
- attributes set in the `@pyi_base` flow-level decorator. Hence,
722
- you can use `@pypi_base` to set packages required by all
723
- steps and use `@pypi` to specify step-specific overrides.
618
+ This decorator is useful if this step may hang indefinitely.
619
+
620
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
621
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
622
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
623
+
624
+ Note that all the values specified in parameters are added together so if you specify
625
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
724
626
 
725
627
  Parameters
726
628
  ----------
727
- packages : Dict[str, str], default: {}
728
- Packages to use for this step. The key is the name of the package
729
- and the value is the version to use.
730
- python : str, optional, default: None
731
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
732
- that the version used will correspond to the version of the Python interpreter used to start the run.
629
+ seconds : int, default 0
630
+ Number of seconds to wait prior to timing out.
631
+ minutes : int, default 0
632
+ Number of minutes to wait prior to timing out.
633
+ hours : int, default 0
634
+ Number of hours to wait prior to timing out.
635
+ """
636
+ ...
637
+
638
+ @typing.overload
639
+ def model(*, load: typing.Union[typing.List[str], str, typing.List[typing.Tuple[str, typing.Optional[str]]]] = None, temp_dir_root: 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]]]:
640
+ """
641
+ Enables loading / saving of models within a step.
642
+
643
+
644
+ Parameters
645
+ ----------
646
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
647
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
648
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
649
+ - `current.checkpoint`
650
+ - `current.model`
651
+ - `current.huggingface_hub`
652
+
653
+ If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
654
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
655
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
656
+
657
+ temp_dir_root : str, default: None
658
+ The root directory under which `current.model.loaded` will store loaded models
659
+
660
+
661
+
662
+ """
663
+ ...
664
+
665
+ @typing.overload
666
+ def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
667
+ ...
668
+
669
+ @typing.overload
670
+ def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
671
+ ...
672
+
673
+ def model(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, load: typing.Union[typing.List[str], str, typing.List[typing.Tuple[str, typing.Optional[str]]]] = None, temp_dir_root: str = None):
674
+ """
675
+ Enables loading / saving of models within a step.
676
+
677
+
678
+ Parameters
679
+ ----------
680
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
681
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
682
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
683
+ - `current.checkpoint`
684
+ - `current.model`
685
+ - `current.huggingface_hub`
686
+
687
+ If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
688
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
689
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
690
+
691
+ temp_dir_root : str, default: None
692
+ The root directory under which `current.model.loaded` will store loaded models
693
+
694
+
695
+
733
696
  """
734
697
  ...
735
698
 
@@ -811,80 +774,59 @@ def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None]
811
774
  ...
812
775
 
813
776
  @typing.overload
814
- def model(*, load: typing.Union[typing.List[str], str, typing.List[typing.Tuple[str, typing.Optional[str]]]] = None, temp_dir_root: 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]]]:
777
+ def conda(*, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
815
778
  """
816
- Enables loading / saving of models within a step.
817
-
779
+ Specifies the Conda environment for the step.
780
+
781
+ Information in this decorator will augment any
782
+ attributes set in the `@conda_base` flow-level decorator. Hence,
783
+ you can use `@conda_base` to set packages required by all
784
+ steps and use `@conda` to specify step-specific overrides.
818
785
 
819
786
  Parameters
820
787
  ----------
821
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
822
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
823
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
824
- - `current.checkpoint`
825
- - `current.model`
826
- - `current.huggingface_hub`
827
-
828
- If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
829
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
830
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
831
-
832
- temp_dir_root : str, default: None
833
- The root directory under which `current.model.loaded` will store loaded models
834
-
835
-
836
-
788
+ packages : Dict[str, str], default {}
789
+ Packages to use for this step. The key is the name of the package
790
+ and the value is the version to use.
791
+ libraries : Dict[str, str], default {}
792
+ Supported for backward compatibility. When used with packages, packages will take precedence.
793
+ python : str, optional, default None
794
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
795
+ that the version used will correspond to the version of the Python interpreter used to start the run.
796
+ disabled : bool, default False
797
+ If set to True, disables @conda.
837
798
  """
838
799
  ...
839
800
 
840
801
  @typing.overload
841
- def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
802
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
842
803
  ...
843
804
 
844
805
  @typing.overload
845
- def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
806
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
846
807
  ...
847
808
 
848
- def model(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, load: typing.Union[typing.List[str], str, typing.List[typing.Tuple[str, typing.Optional[str]]]] = None, temp_dir_root: str = None):
809
+ def conda(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
849
810
  """
850
- Enables loading / saving of models within a step.
811
+ Specifies the Conda environment for the step.
851
812
 
813
+ Information in this decorator will augment any
814
+ attributes set in the `@conda_base` flow-level decorator. Hence,
815
+ you can use `@conda_base` to set packages required by all
816
+ steps and use `@conda` to specify step-specific overrides.
852
817
 
853
818
  Parameters
854
819
  ----------
855
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
856
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
857
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
858
- - `current.checkpoint`
859
- - `current.model`
860
- - `current.huggingface_hub`
861
-
862
- If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
863
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
864
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
865
-
866
- temp_dir_root : str, default: None
867
- The root directory under which `current.model.loaded` will store loaded models
868
-
869
-
870
-
871
- """
872
- ...
873
-
874
- @typing.overload
875
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
876
- """
877
- Internal decorator to support Fast bakery
878
- """
879
- ...
880
-
881
- @typing.overload
882
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
883
- ...
884
-
885
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
886
- """
887
- Internal decorator to support Fast bakery
820
+ packages : Dict[str, str], default {}
821
+ Packages to use for this step. The key is the name of the package
822
+ and the value is the version to use.
823
+ libraries : Dict[str, str], default {}
824
+ Supported for backward compatibility. When used with packages, packages will take precedence.
825
+ python : str, optional, default None
826
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
827
+ that the version used will correspond to the version of the Python interpreter used to start the run.
828
+ disabled : bool, default False
829
+ If set to True, disables @conda.
888
830
  """
889
831
  ...
890
832
 
@@ -922,109 +864,167 @@ def secrets(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None],
922
864
  ...
923
865
 
924
866
  @typing.overload
925
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
867
+ 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]]]:
926
868
  """
927
- Decorator prototype for all step decorators. This function gets specialized
928
- and imported for all decorators types by _import_plugin_decorators().
869
+ Specifies environment variables to be set prior to the execution of a step.
870
+
871
+ Parameters
872
+ ----------
873
+ vars : Dict[str, str], default {}
874
+ Dictionary of environment variables to set.
929
875
  """
930
876
  ...
931
877
 
932
878
  @typing.overload
933
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
879
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
934
880
  ...
935
881
 
936
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
882
+ @typing.overload
883
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
884
+ ...
885
+
886
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
937
887
  """
938
- Decorator prototype for all step decorators. This function gets specialized
939
- and imported for all decorators types by _import_plugin_decorators().
888
+ Specifies environment variables to be set prior to the execution of a step.
889
+
890
+ Parameters
891
+ ----------
892
+ vars : Dict[str, str], default {}
893
+ Dictionary of environment variables to set.
940
894
  """
941
895
  ...
942
896
 
943
897
  @typing.overload
944
- def conda(*, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
898
+ 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]]]:
945
899
  """
946
- Specifies the Conda environment for the step.
900
+ Specifies the PyPI packages for the step.
947
901
 
948
902
  Information in this decorator will augment any
949
- attributes set in the `@conda_base` flow-level decorator. Hence,
950
- you can use `@conda_base` to set packages required by all
951
- steps and use `@conda` to specify step-specific overrides.
903
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
904
+ you can use `@pypi_base` to set packages required by all
905
+ steps and use `@pypi` to specify step-specific overrides.
952
906
 
953
907
  Parameters
954
908
  ----------
955
- packages : Dict[str, str], default {}
909
+ packages : Dict[str, str], default: {}
956
910
  Packages to use for this step. The key is the name of the package
957
911
  and the value is the version to use.
958
- libraries : Dict[str, str], default {}
959
- Supported for backward compatibility. When used with packages, packages will take precedence.
960
- python : str, optional, default None
912
+ python : str, optional, default: None
961
913
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
962
914
  that the version used will correspond to the version of the Python interpreter used to start the run.
963
- disabled : bool, default False
964
- If set to True, disables @conda.
965
915
  """
966
916
  ...
967
917
 
968
918
  @typing.overload
969
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
919
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
970
920
  ...
971
921
 
972
922
  @typing.overload
973
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
923
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
974
924
  ...
975
925
 
976
- def conda(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
926
+ 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):
977
927
  """
978
- Specifies the Conda environment for the step.
928
+ Specifies the PyPI packages for the step.
979
929
 
980
930
  Information in this decorator will augment any
981
- attributes set in the `@conda_base` flow-level decorator. Hence,
982
- you can use `@conda_base` to set packages required by all
983
- steps and use `@conda` to specify step-specific overrides.
931
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
932
+ you can use `@pypi_base` to set packages required by all
933
+ steps and use `@pypi` to specify step-specific overrides.
984
934
 
985
935
  Parameters
986
936
  ----------
987
- packages : Dict[str, str], default {}
937
+ packages : Dict[str, str], default: {}
988
938
  Packages to use for this step. The key is the name of the package
989
939
  and the value is the version to use.
990
- libraries : Dict[str, str], default {}
991
- Supported for backward compatibility. When used with packages, packages will take precedence.
992
- python : str, optional, default None
940
+ python : str, optional, default: None
993
941
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
994
942
  that the version used will correspond to the version of the Python interpreter used to start the run.
995
- disabled : bool, default False
996
- If set to True, disables @conda.
997
943
  """
998
944
  ...
999
945
 
1000
946
  @typing.overload
1001
- 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]]]:
947
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1002
948
  """
1003
- Specifies environment variables to be set prior to the execution of a step.
1004
-
1005
- Parameters
1006
- ----------
1007
- vars : Dict[str, str], default {}
1008
- Dictionary of environment variables to set.
949
+ Decorator prototype for all step decorators. This function gets specialized
950
+ and imported for all decorators types by _import_plugin_decorators().
1009
951
  """
1010
952
  ...
1011
953
 
1012
954
  @typing.overload
1013
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
955
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1014
956
  ...
1015
957
 
1016
- @typing.overload
1017
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
958
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
959
+ """
960
+ Decorator prototype for all step decorators. This function gets specialized
961
+ and imported for all decorators types by _import_plugin_decorators().
962
+ """
1018
963
  ...
1019
964
 
1020
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
965
+ 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, node_selector: typing.Union[typing.Dict[str, str], str, None] = None, namespace: str = "METAFLOW_KUBERNETES_NAMESPACE", gpu: typing.Optional[int] = None, gpu_vendor: str = "KUBERNETES_GPU_VENDOR", tolerations: typing.List[str] = [], use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = "/metaflow_temp", persistent_volume_claims: typing.Optional[typing.Dict[str, str]] = None, shared_memory: typing.Optional[int] = None, port: typing.Optional[int] = None, compute_pool: typing.Optional[str] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
1021
966
  """
1022
- Specifies environment variables to be set prior to the execution of a step.
967
+ Specifies that this step should execute on Kubernetes.
1023
968
 
1024
969
  Parameters
1025
970
  ----------
1026
- vars : Dict[str, str], default {}
1027
- Dictionary of environment variables to set.
971
+ cpu : int, default 1
972
+ Number of CPUs required for this step. If `@resources` is
973
+ also present, the maximum value from all decorators is used.
974
+ memory : int, default 4096
975
+ Memory size (in MB) required for this step. If
976
+ `@resources` is also present, the maximum value from all decorators is
977
+ used.
978
+ disk : int, default 10240
979
+ Disk size (in MB) required for this step. If
980
+ `@resources` is also present, the maximum value from all decorators is
981
+ used.
982
+ image : str, optional, default None
983
+ Docker image to use when launching on Kubernetes. If not specified, and
984
+ METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
985
+ not, a default Docker image mapping to the current version of Python is used.
986
+ image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
987
+ If given, the imagePullPolicy to be applied to the Docker image of the step.
988
+ service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
989
+ Kubernetes service account to use when launching pod in Kubernetes.
990
+ secrets : List[str], optional, default None
991
+ Kubernetes secrets to use when launching pod in Kubernetes. These
992
+ secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
993
+ in Metaflow configuration.
994
+ node_selector: Union[Dict[str,str], str], optional, default None
995
+ Kubernetes node selector(s) to apply to the pod running the task.
996
+ Can be passed in as a comma separated string of values e.g. "kubernetes.io/os=linux,kubernetes.io/arch=amd64"
997
+ or as a dictionary {"kubernetes.io/os": "linux", "kubernetes.io/arch": "amd64"}
998
+ namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
999
+ Kubernetes namespace to use when launching pod in Kubernetes.
1000
+ gpu : int, optional, default None
1001
+ Number of GPUs required for this step. A value of zero implies that
1002
+ the scheduled node should not have GPUs.
1003
+ gpu_vendor : str, default KUBERNETES_GPU_VENDOR
1004
+ The vendor of the GPUs to be used for this step.
1005
+ tolerations : List[str], default []
1006
+ The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
1007
+ Kubernetes tolerations to use when launching pod in Kubernetes.
1008
+ use_tmpfs : bool, default False
1009
+ This enables an explicit tmpfs mount for this step.
1010
+ tmpfs_tempdir : bool, default True
1011
+ sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
1012
+ tmpfs_size : int, optional, default: None
1013
+ The value for the size (in MiB) of the tmpfs mount for this step.
1014
+ This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
1015
+ memory allocated for this step.
1016
+ tmpfs_path : str, optional, default /metaflow_temp
1017
+ Path to tmpfs mount for this step.
1018
+ persistent_volume_claims : Dict[str, str], optional, default None
1019
+ A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
1020
+ volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
1021
+ shared_memory: int, optional
1022
+ Shared memory size (in MiB) required for this step
1023
+ port: int, optional
1024
+ Port number to specify in the Kubernetes job object
1025
+ compute_pool : str, optional, default None
1026
+ Compute pool to be used for for this step.
1027
+ If not specified, any accessible compute pool within the perimeter is used.
1028
1028
  """
1029
1029
  ...
1030
1030
 
@@ -1089,59 +1089,6 @@ def checkpoint(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None
1089
1089
  """
1090
1090
  ...
1091
1091
 
1092
- @typing.overload
1093
- 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]]]:
1094
- """
1095
- Specifies the number of times the task corresponding
1096
- to a step needs to be retried.
1097
-
1098
- This decorator is useful for handling transient errors, such as networking issues.
1099
- If your task contains operations that can't be retried safely, e.g. database updates,
1100
- it is advisable to annotate it with `@retry(times=0)`.
1101
-
1102
- This can be used in conjunction with the `@catch` decorator. The `@catch`
1103
- decorator will execute a no-op task after all retries have been exhausted,
1104
- ensuring that the flow execution can continue.
1105
-
1106
- Parameters
1107
- ----------
1108
- times : int, default 3
1109
- Number of times to retry this task.
1110
- minutes_between_retries : int, default 2
1111
- Number of minutes between retries.
1112
- """
1113
- ...
1114
-
1115
- @typing.overload
1116
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1117
- ...
1118
-
1119
- @typing.overload
1120
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1121
- ...
1122
-
1123
- 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):
1124
- """
1125
- Specifies the number of times the task corresponding
1126
- to a step needs to be retried.
1127
-
1128
- This decorator is useful for handling transient errors, such as networking issues.
1129
- If your task contains operations that can't be retried safely, e.g. database updates,
1130
- it is advisable to annotate it with `@retry(times=0)`.
1131
-
1132
- This can be used in conjunction with the `@catch` decorator. The `@catch`
1133
- decorator will execute a no-op task after all retries have been exhausted,
1134
- ensuring that the flow execution can continue.
1135
-
1136
- Parameters
1137
- ----------
1138
- times : int, default 3
1139
- Number of times to retry this task.
1140
- minutes_between_retries : int, default 2
1141
- Number of minutes between retries.
1142
- """
1143
- ...
1144
-
1145
1092
  @typing.overload
1146
1093
  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]]]:
1147
1094
  """
@@ -1194,41 +1141,55 @@ def card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typ
1194
1141
  ...
1195
1142
 
1196
1143
  @typing.overload
1197
- def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1144
+ 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]]]:
1198
1145
  """
1199
- Specifies the PyPI packages for all steps of the flow.
1146
+ Specifies the number of times the task corresponding
1147
+ to a step needs to be retried.
1148
+
1149
+ This decorator is useful for handling transient errors, such as networking issues.
1150
+ If your task contains operations that can't be retried safely, e.g. database updates,
1151
+ it is advisable to annotate it with `@retry(times=0)`.
1152
+
1153
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
1154
+ decorator will execute a no-op task after all retries have been exhausted,
1155
+ ensuring that the flow execution can continue.
1200
1156
 
1201
- Use `@pypi_base` to set common packages required by all
1202
- steps and use `@pypi` to specify step-specific overrides.
1203
1157
  Parameters
1204
1158
  ----------
1205
- packages : Dict[str, str], default: {}
1206
- Packages to use for this flow. The key is the name of the package
1207
- and the value is the version to use.
1208
- python : str, optional, default: None
1209
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1210
- that the version used will correspond to the version of the Python interpreter used to start the run.
1159
+ times : int, default 3
1160
+ Number of times to retry this task.
1161
+ minutes_between_retries : int, default 2
1162
+ Number of minutes between retries.
1211
1163
  """
1212
1164
  ...
1213
1165
 
1214
1166
  @typing.overload
1215
- def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1167
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1216
1168
  ...
1217
1169
 
1218
- def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1170
+ @typing.overload
1171
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1172
+ ...
1173
+
1174
+ 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):
1219
1175
  """
1220
- Specifies the PyPI packages for all steps of the flow.
1176
+ Specifies the number of times the task corresponding
1177
+ to a step needs to be retried.
1178
+
1179
+ This decorator is useful for handling transient errors, such as networking issues.
1180
+ If your task contains operations that can't be retried safely, e.g. database updates,
1181
+ it is advisable to annotate it with `@retry(times=0)`.
1182
+
1183
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
1184
+ decorator will execute a no-op task after all retries have been exhausted,
1185
+ ensuring that the flow execution can continue.
1221
1186
 
1222
- Use `@pypi_base` to set common packages required by all
1223
- steps and use `@pypi` to specify step-specific overrides.
1224
1187
  Parameters
1225
1188
  ----------
1226
- packages : Dict[str, str], default: {}
1227
- Packages to use for this flow. The key is the name of the package
1228
- and the value is the version to use.
1229
- python : str, optional, default: None
1230
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1231
- that the version used will correspond to the version of the Python interpreter used to start the run.
1189
+ times : int, default 3
1190
+ Number of times to retry this task.
1191
+ minutes_between_retries : int, default 2
1192
+ Number of minutes between retries.
1232
1193
  """
1233
1194
  ...
1234
1195
 
@@ -1297,132 +1258,37 @@ def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *
1297
1258
  ```
1298
1259
  @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1299
1260
  ```
1300
- This decorator respects the @project decorator and triggers the flow
1301
- when upstream runs within the same namespace complete successfully
1302
-
1303
- Additionally, you can specify project aware upstream flow dependencies
1304
- by specifying the fully qualified project_flow_name.
1305
- ```
1306
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1307
- ```
1308
- or
1309
- ```
1310
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1311
- ```
1312
-
1313
- You can also specify just the project or project branch (other values will be
1314
- inferred from the current project or project branch):
1315
- ```
1316
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1317
- ```
1318
-
1319
- Note that `branch` is typically one of:
1320
- - `prod`
1321
- - `user.bob`
1322
- - `test.my_experiment`
1323
- - `prod.staging`
1324
-
1325
- Parameters
1326
- ----------
1327
- flow : Union[str, Dict[str, str]], optional, default None
1328
- Upstream flow dependency for this flow.
1329
- flows : List[Union[str, Dict[str, str]]], default []
1330
- Upstream flow dependencies for this flow.
1331
- options : Dict[str, Any], default {}
1332
- Backend-specific configuration for tuning eventing behavior.
1333
-
1334
-
1335
- """
1336
- ...
1337
-
1338
- @typing.overload
1339
- def trigger(*, event: typing.Union[str, typing.Dict[str, typing.Any], None] = None, events: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], options: typing.Dict[str, typing.Any] = {}) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1340
- """
1341
- Specifies the event(s) that this flow depends on.
1342
-
1343
- ```
1344
- @trigger(event='foo')
1345
- ```
1346
- or
1347
- ```
1348
- @trigger(events=['foo', 'bar'])
1349
- ```
1350
-
1351
- Additionally, you can specify the parameter mappings
1352
- to map event payload to Metaflow parameters for the flow.
1353
- ```
1354
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1355
- ```
1356
- or
1357
- ```
1358
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1359
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1360
- ```
1361
-
1362
- 'parameters' can also be a list of strings and tuples like so:
1363
- ```
1364
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1365
- ```
1366
- This is equivalent to:
1367
- ```
1368
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1369
- ```
1370
-
1371
- Parameters
1372
- ----------
1373
- event : Union[str, Dict[str, Any]], optional, default None
1374
- Event dependency for this flow.
1375
- events : List[Union[str, Dict[str, Any]]], default []
1376
- Events dependency for this flow.
1377
- options : Dict[str, Any], default {}
1378
- Backend-specific configuration for tuning eventing behavior.
1379
-
1380
-
1381
- """
1382
- ...
1383
-
1384
- @typing.overload
1385
- def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1386
- ...
1387
-
1388
- def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: typing.Union[str, typing.Dict[str, typing.Any], None] = None, events: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], options: typing.Dict[str, typing.Any] = {}):
1389
- """
1390
- Specifies the event(s) that this flow depends on.
1391
-
1392
- ```
1393
- @trigger(event='foo')
1394
- ```
1395
- or
1396
- ```
1397
- @trigger(events=['foo', 'bar'])
1398
- ```
1261
+ This decorator respects the @project decorator and triggers the flow
1262
+ when upstream runs within the same namespace complete successfully
1399
1263
 
1400
- Additionally, you can specify the parameter mappings
1401
- to map event payload to Metaflow parameters for the flow.
1264
+ Additionally, you can specify project aware upstream flow dependencies
1265
+ by specifying the fully qualified project_flow_name.
1402
1266
  ```
1403
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1267
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1404
1268
  ```
1405
1269
  or
1406
1270
  ```
1407
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1408
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1271
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1409
1272
  ```
1410
1273
 
1411
- 'parameters' can also be a list of strings and tuples like so:
1412
- ```
1413
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1414
- ```
1415
- This is equivalent to:
1274
+ You can also specify just the project or project branch (other values will be
1275
+ inferred from the current project or project branch):
1416
1276
  ```
1417
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1277
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1418
1278
  ```
1419
1279
 
1280
+ Note that `branch` is typically one of:
1281
+ - `prod`
1282
+ - `user.bob`
1283
+ - `test.my_experiment`
1284
+ - `prod.staging`
1285
+
1420
1286
  Parameters
1421
1287
  ----------
1422
- event : Union[str, Dict[str, Any]], optional, default None
1423
- Event dependency for this flow.
1424
- events : List[Union[str, Dict[str, Any]]], default []
1425
- Events dependency for this flow.
1288
+ flow : Union[str, Dict[str, str]], optional, default None
1289
+ Upstream flow dependency for this flow.
1290
+ flows : List[Union[str, Dict[str, str]]], default []
1291
+ Upstream flow dependencies for this flow.
1426
1292
  options : Dict[str, Any], default {}
1427
1293
  Backend-specific configuration for tuning eventing behavior.
1428
1294
 
@@ -1430,6 +1296,48 @@ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: t
1430
1296
  """
1431
1297
  ...
1432
1298
 
1299
+ 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]]:
1300
+ """
1301
+ 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)
1302
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1303
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1304
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
1305
+ starts only after all sensors finish.
1306
+
1307
+ Parameters
1308
+ ----------
1309
+ timeout : int
1310
+ Time, in seconds before the task times out and fails. (Default: 3600)
1311
+ poke_interval : int
1312
+ Time in seconds that the job should wait in between each try. (Default: 60)
1313
+ mode : str
1314
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1315
+ exponential_backoff : bool
1316
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1317
+ pool : str
1318
+ the slot pool this task should run in,
1319
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1320
+ soft_fail : bool
1321
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1322
+ name : str
1323
+ Name of the sensor on Airflow
1324
+ description : str
1325
+ Description of sensor in the Airflow UI
1326
+ bucket_key : Union[str, List[str]]
1327
+ The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1328
+ When it's specified as a full s3:// url, please leave `bucket_name` as None
1329
+ bucket_name : str
1330
+ Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1331
+ When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1332
+ wildcard_match : bool
1333
+ whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1334
+ aws_conn_id : str
1335
+ a reference to the s3 connection on Airflow. (Default: None)
1336
+ verify : bool
1337
+ Whether or not to verify SSL certificates for S3 connection. (Default: None)
1338
+ """
1339
+ ...
1340
+
1433
1341
  def nim(*, models: "list[NIM]", backend: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1434
1342
  """
1435
1343
  This decorator is used to run NIM containers in Metaflow tasks as sidecars.
@@ -1528,6 +1436,45 @@ def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly:
1528
1436
  """
1529
1437
  ...
1530
1438
 
1439
+ @typing.overload
1440
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1441
+ """
1442
+ Specifies the PyPI packages for all steps of the flow.
1443
+
1444
+ Use `@pypi_base` to set common packages required by all
1445
+ steps and use `@pypi` to specify step-specific overrides.
1446
+ Parameters
1447
+ ----------
1448
+ packages : Dict[str, str], default: {}
1449
+ Packages to use for this flow. The key is the name of the package
1450
+ and the value is the version to use.
1451
+ python : str, optional, default: None
1452
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1453
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1454
+ """
1455
+ ...
1456
+
1457
+ @typing.overload
1458
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1459
+ ...
1460
+
1461
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1462
+ """
1463
+ Specifies the PyPI packages for all steps of the flow.
1464
+
1465
+ Use `@pypi_base` to set common packages required by all
1466
+ steps and use `@pypi` to specify step-specific overrides.
1467
+ Parameters
1468
+ ----------
1469
+ packages : Dict[str, str], default: {}
1470
+ Packages to use for this flow. The key is the name of the package
1471
+ and the value is the version to use.
1472
+ python : str, optional, default: None
1473
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1474
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1475
+ """
1476
+ ...
1477
+
1531
1478
  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]]:
1532
1479
  """
1533
1480
  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.
@@ -1570,48 +1517,6 @@ def airflow_external_task_sensor(*, timeout: int, poke_interval: int, mode: str,
1570
1517
  """
1571
1518
  ...
1572
1519
 
1573
- 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]]:
1574
- """
1575
- 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)
1576
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1577
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1578
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
1579
- starts only after all sensors finish.
1580
-
1581
- Parameters
1582
- ----------
1583
- timeout : int
1584
- Time, in seconds before the task times out and fails. (Default: 3600)
1585
- poke_interval : int
1586
- Time in seconds that the job should wait in between each try. (Default: 60)
1587
- mode : str
1588
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1589
- exponential_backoff : bool
1590
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1591
- pool : str
1592
- the slot pool this task should run in,
1593
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1594
- soft_fail : bool
1595
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1596
- name : str
1597
- Name of the sensor on Airflow
1598
- description : str
1599
- Description of sensor in the Airflow UI
1600
- bucket_key : Union[str, List[str]]
1601
- The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1602
- When it's specified as a full s3:// url, please leave `bucket_name` as None
1603
- bucket_name : str
1604
- Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1605
- When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1606
- wildcard_match : bool
1607
- whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1608
- aws_conn_id : str
1609
- a reference to the s3 connection on Airflow. (Default: None)
1610
- verify : bool
1611
- Whether or not to verify SSL certificates for S3 connection. (Default: None)
1612
- """
1613
- ...
1614
-
1615
1520
  @typing.overload
1616
1521
  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]]:
1617
1522
  """
@@ -1661,6 +1566,101 @@ def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packa
1661
1566
  """
1662
1567
  ...
1663
1568
 
1569
+ @typing.overload
1570
+ def trigger(*, event: typing.Union[str, typing.Dict[str, typing.Any], None] = None, events: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], options: typing.Dict[str, typing.Any] = {}) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1571
+ """
1572
+ Specifies the event(s) that this flow depends on.
1573
+
1574
+ ```
1575
+ @trigger(event='foo')
1576
+ ```
1577
+ or
1578
+ ```
1579
+ @trigger(events=['foo', 'bar'])
1580
+ ```
1581
+
1582
+ Additionally, you can specify the parameter mappings
1583
+ to map event payload to Metaflow parameters for the flow.
1584
+ ```
1585
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1586
+ ```
1587
+ or
1588
+ ```
1589
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1590
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1591
+ ```
1592
+
1593
+ 'parameters' can also be a list of strings and tuples like so:
1594
+ ```
1595
+ @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1596
+ ```
1597
+ This is equivalent to:
1598
+ ```
1599
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1600
+ ```
1601
+
1602
+ Parameters
1603
+ ----------
1604
+ event : Union[str, Dict[str, Any]], optional, default None
1605
+ Event dependency for this flow.
1606
+ events : List[Union[str, Dict[str, Any]]], default []
1607
+ Events dependency for this flow.
1608
+ options : Dict[str, Any], default {}
1609
+ Backend-specific configuration for tuning eventing behavior.
1610
+
1611
+
1612
+ """
1613
+ ...
1614
+
1615
+ @typing.overload
1616
+ def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1617
+ ...
1618
+
1619
+ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: typing.Union[str, typing.Dict[str, typing.Any], None] = None, events: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], options: typing.Dict[str, typing.Any] = {}):
1620
+ """
1621
+ Specifies the event(s) that this flow depends on.
1622
+
1623
+ ```
1624
+ @trigger(event='foo')
1625
+ ```
1626
+ or
1627
+ ```
1628
+ @trigger(events=['foo', 'bar'])
1629
+ ```
1630
+
1631
+ Additionally, you can specify the parameter mappings
1632
+ to map event payload to Metaflow parameters for the flow.
1633
+ ```
1634
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1635
+ ```
1636
+ or
1637
+ ```
1638
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1639
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1640
+ ```
1641
+
1642
+ 'parameters' can also be a list of strings and tuples like so:
1643
+ ```
1644
+ @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1645
+ ```
1646
+ This is equivalent to:
1647
+ ```
1648
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1649
+ ```
1650
+
1651
+ Parameters
1652
+ ----------
1653
+ event : Union[str, Dict[str, Any]], optional, default None
1654
+ Event dependency for this flow.
1655
+ events : List[Union[str, Dict[str, Any]]], default []
1656
+ Events dependency for this flow.
1657
+ options : Dict[str, Any], default {}
1658
+ Backend-specific configuration for tuning eventing behavior.
1659
+
1660
+
1661
+ """
1662
+ ...
1663
+
1664
1664
  def namespace(ns: typing.Optional[str]) -> typing.Optional[str]:
1665
1665
  """
1666
1666
  Switch namespace to the one provided.
@@ -3107,20 +3107,35 @@ class Checkpoint(object, metaclass=type):
3107
3107
  @property
3108
3108
  def directory(self):
3109
3109
  ...
3110
- def save(self, path = None, metadata = None, latest = True, name = "mfchckpt"):
3110
+ def save(self, path = None, metadata = None, latest = True, name = "mfchckpt", storage_format = "files"):
3111
3111
  """
3112
- saves the checkpoint to the datastore (resemble's a create op)
3112
+ Saves the checkpoint to the datastore
3113
3113
 
3114
3114
  Parameters
3115
3115
  ----------
3116
+ path : Optional[Union[str, os.PathLike]], default: None
3117
+ The path to save the checkpoint. Accepts a file path or a directory path.
3118
+ - If a directory path is provided, all the contents within that directory will be saved.
3119
+ When a checkpoint is reloaded during task retries, `the current.checkpoint.directory` will
3120
+ contain the contents of this directory.
3121
+ - If a file path is provided, the file will be directly saved to the datastore (with the same filename).
3122
+ When the checkpoint is reloaded during task retries, the file with the same name will be available in the
3123
+ `current.checkpoint.directory`.
3124
+ - If no path is provided then the `Checkpoint.directory` will be saved as the checkpoint.
3116
3125
 
3117
- - `path` (str or os.PathLike):
3118
- - path to directory
3119
- - path to file
3120
- - `name`:
3121
- - name of the checkpoint
3122
- - `metadata`:
3123
- - metadata of the checkpoint
3126
+ name : Optional[str], default: "mfchckpt"
3127
+ The name of the checkpoint.
3128
+
3129
+ metadata : Optional[Dict], default: {}
3130
+ Any metadata that needs to be saved with the checkpoint.
3131
+
3132
+ latest : bool, default: True
3133
+ If True, the checkpoint will be marked as the latest checkpoint.
3134
+ This helps determine if the checkpoint gets loaded when the task restarts.
3135
+
3136
+ storage_format : str, default: files
3137
+ If `tar`, the contents of the directory will be tarred before saving to the datastore.
3138
+ If `files`, saves directory directly to the datastore.
3124
3139
  """
3125
3140
  ...
3126
3141
  def __enter__(self):
@@ -3138,7 +3153,6 @@ class Checkpoint(object, metaclass=type):
3138
3153
  ```python
3139
3154
 
3140
3155
  Checkpoint().list(name="best") # lists checkpoints in the current task with the name "best"
3141
- Checkpoint().list(name="best", task="anotherflow/somerunid/somestep/sometask") # lists checkpoints in the another task with the name "best"
3142
3156
  Checkpoint().list(task="anotherflow/somerunid/somestep/sometask", name="best") # Identical as the above one but
3143
3157
  Checkpoint().list() # lists all the checkpoints in the current task
3144
3158