ob-metaflow-stubs 6.0.3.107__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 +583 -569
  2. metaflow-stubs/cards.pyi +5 -5
  3. metaflow-stubs/cli.pyi +3 -3
  4. metaflow-stubs/client/__init__.pyi +3 -3
  5. metaflow-stubs/client/core.pyi +6 -6
  6. metaflow-stubs/client/filecache.pyi +3 -3
  7. metaflow-stubs/clone_util.pyi +2 -2
  8. metaflow-stubs/events.pyi +2 -2
  9. metaflow-stubs/exception.pyi +2 -2
  10. metaflow-stubs/flowspec.pyi +5 -5
  11. metaflow-stubs/generated_for.txt +1 -1
  12. metaflow-stubs/includefile.pyi +5 -5
  13. metaflow-stubs/info_file.pyi +2 -2
  14. metaflow-stubs/metadata/metadata.pyi +3 -3
  15. metaflow-stubs/metadata/util.pyi +2 -2
  16. metaflow-stubs/metaflow_config.pyi +2 -2
  17. metaflow-stubs/metaflow_current.pyi +67 -67
  18. metaflow-stubs/mflog/mflog.pyi +2 -2
  19. metaflow-stubs/multicore_utils.pyi +2 -2
  20. metaflow-stubs/parameters.pyi +3 -3
  21. metaflow-stubs/plugins/__init__.pyi +3 -3
  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 +2 -2
  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 +3 -3
  33. metaflow-stubs/plugins/argo/argo_workflows_cli.pyi +6 -6
  34. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +5 -5
  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 +2 -2
  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 +2 -2
  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 +2 -2
  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 +3 -3
  67. metaflow-stubs/plugins/cards/card_modules/card.pyi +2 -2
  68. metaflow-stubs/plugins/cards/card_modules/components.pyi +3 -3
  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 +2 -2
  72. metaflow-stubs/plugins/cards/card_resolver.pyi +2 -2
  73. metaflow-stubs/plugins/cards/component_serializer.pyi +3 -3
  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 +6 -6
  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 +3 -3
  88. metaflow-stubs/plugins/gcp/__init__.pyi +2 -2
  89. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +2 -2
  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 +3 -3
  97. metaflow-stubs/plugins/kubernetes/kubernetes_cli.pyi +3 -3
  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 +3 -3
  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 +3 -3
  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 +4 -4
  119. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +3 -3
  120. metaflow-stubs/plugins/timeout_decorator.pyi +2 -2
  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 +5 -5
  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 +2 -2
  134. metaflow-stubs/tagging_util.pyi +2 -2
  135. metaflow-stubs/tuple_util.pyi +2 -2
  136. {ob_metaflow_stubs-6.0.3.107.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.107.dist-info/RECORD +0 -140
  139. {ob_metaflow_stubs-6.0.3.107.dist-info → ob_metaflow_stubs-6.0.3.109rc0.dist-info}/WHEEL +0 -0
  140. {ob_metaflow_stubs-6.0.3.107.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.24.1+obcheckpoint(0.0.11);ob(v1) #
4
- # Generated on 2024-10-07T18:39:02.357763 #
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.metaflow_current
12
- import metaflow._vendor.click.types
13
- import metaflow.client.core
14
- import datetime
11
+ import metaflow.events
15
12
  import metaflow_extensions.obcheckpoint.plugins.machine_learning_utilities.datastructures
13
+ import metaflow.runner.metaflow_runner
14
+ import datetime
15
+ import metaflow.metaflow_current
16
16
  import typing
17
- import metaflow.flowspec
17
+ import metaflow._vendor.click.types
18
18
  import metaflow.parameters
19
+ import metaflow.client.core
19
20
  import metaflow.datastore.inputs
20
- import metaflow.runner.metaflow_runner
21
- import metaflow.events
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
 
@@ -512,6 +512,23 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
512
512
  """
513
513
  ...
514
514
 
515
+ @typing.overload
516
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
517
+ """
518
+ Internal decorator to support Fast bakery
519
+ """
520
+ ...
521
+
522
+ @typing.overload
523
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
524
+ ...
525
+
526
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
527
+ """
528
+ Internal decorator to support Fast bakery
529
+ """
530
+ ...
531
+
515
532
  @typing.overload
516
533
  def catch(*, var: typing.Optional[str] = None, print_exception: bool = True) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
517
534
  """
@@ -562,92 +579,83 @@ def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
562
579
  ...
563
580
 
564
581
  @typing.overload
565
- 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]]]:
566
583
  """
567
- Specifies the PyPI packages for the step.
584
+ Specifies a timeout for your step.
568
585
 
569
- Information in this decorator will augment any
570
- attributes set in the `@pyi_base` flow-level decorator. Hence,
571
- you can use `@pypi_base` to set packages required by all
572
- 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.
573
594
 
574
595
  Parameters
575
596
  ----------
576
- packages : Dict[str, str], default: {}
577
- Packages to use for this step. The key is the name of the package
578
- and the value is the version to use.
579
- python : str, optional, default: None
580
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
581
- 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.
582
603
  """
583
604
  ...
584
605
 
585
606
  @typing.overload
586
- 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]:
587
608
  ...
588
609
 
589
610
  @typing.overload
590
- 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]:
591
612
  ...
592
613
 
593
- 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):
594
615
  """
595
- Specifies the PyPI packages for the step.
616
+ Specifies a timeout for your step.
596
617
 
597
- Information in this decorator will augment any
598
- attributes set in the `@pyi_base` flow-level decorator. Hence,
599
- you can use `@pypi_base` to set packages required by all
600
- 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.
601
626
 
602
627
  Parameters
603
628
  ----------
604
- packages : Dict[str, str], default: {}
605
- Packages to use for this step. The key is the name of the package
606
- and the value is the version to use.
607
- python : str, optional, default: None
608
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
609
- that the version used will correspond to the version of the Python interpreter used to start the run.
610
- """
611
- ...
612
-
613
- @typing.overload
614
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
615
- """
616
- Internal decorator to support Fast bakery
617
- """
618
- ...
619
-
620
- @typing.overload
621
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
622
- ...
623
-
624
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
625
- """
626
- Internal decorator to support Fast bakery
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.
627
635
  """
628
636
  ...
629
637
 
630
638
  @typing.overload
631
- def checkpoint(*, load_policy: str = "fresh", 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]]]:
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]]]:
632
640
  """
633
- Enables checkpointing for a step.
641
+ Enables loading / saving of models within a step.
634
642
 
635
643
 
636
644
  Parameters
637
645
  ----------
638
- load_policy : str, default: "fresh"
639
- The policy for loading the checkpoint. The following policies are supported:
640
- - "eager": Loads the the latest available checkpoint within the namespace.
641
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
642
- will be loaded at the start of the task.
643
- - "none": Do not load any checkpoint
644
- - "fresh": Loads the lastest checkpoint created within the running Task.
645
- This mode helps loading checkpoints across various retry attempts of the same task.
646
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
647
- created within the task will be loaded when the task is retries execution on failure.
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.
648
656
 
649
657
  temp_dir_root : str, default: None
650
- The root directory under which `current.checkpoint.directory` will be created.
658
+ The root directory under which `current.model.loaded` will store loaded models
651
659
 
652
660
 
653
661
 
@@ -655,33 +663,33 @@ def checkpoint(*, load_policy: str = "fresh", temp_dir_root: str = None) -> typi
655
663
  ...
656
664
 
657
665
  @typing.overload
658
- def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
666
+ def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
659
667
  ...
660
668
 
661
669
  @typing.overload
662
- def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
670
+ def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
663
671
  ...
664
672
 
665
- def checkpoint(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, load_policy: str = "fresh", temp_dir_root: str = None):
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):
666
674
  """
667
- Enables checkpointing for a step.
675
+ Enables loading / saving of models within a step.
668
676
 
669
677
 
670
678
  Parameters
671
679
  ----------
672
- load_policy : str, default: "fresh"
673
- The policy for loading the checkpoint. The following policies are supported:
674
- - "eager": Loads the the latest available checkpoint within the namespace.
675
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
676
- will be loaded at the start of the task.
677
- - "none": Do not load any checkpoint
678
- - "fresh": Loads the lastest checkpoint created within the running Task.
679
- This mode helps loading checkpoints across various retry attempts of the same task.
680
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
681
- created within the task will be loaded when the task is retries execution on failure.
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.
682
690
 
683
691
  temp_dir_root : str, default: None
684
- The root directory under which `current.checkpoint.directory` will be created.
692
+ The root directory under which `current.model.loaded` will store loaded models
685
693
 
686
694
 
687
695
 
@@ -766,82 +774,59 @@ def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None]
766
774
  ...
767
775
 
768
776
  @typing.overload
769
- 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]]]:
770
778
  """
771
- Enables loading / saving of models within a step.
779
+ Specifies the Conda environment for the step.
772
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.
773
785
 
774
786
  Parameters
775
787
  ----------
776
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
777
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
778
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
779
- - `current.checkpoint`
780
- - `current.model`
781
- - `current.huggingface_hub`
782
-
783
- 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
784
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
785
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
786
-
787
- temp_dir_root : str, default: None
788
- The root directory under which `current.model.loaded` will store loaded models
789
-
790
-
791
-
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.
792
798
  """
793
799
  ...
794
800
 
795
801
  @typing.overload
796
- 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]:
797
803
  ...
798
804
 
799
805
  @typing.overload
800
- 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]:
801
807
  ...
802
808
 
803
- 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):
804
810
  """
805
- Enables loading / saving of models within a step.
811
+ Specifies the Conda environment for the step.
806
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.
807
817
 
808
818
  Parameters
809
819
  ----------
810
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
811
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
812
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
813
- - `current.checkpoint`
814
- - `current.model`
815
- - `current.huggingface_hub`
816
-
817
- 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
818
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
819
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
820
-
821
- temp_dir_root : str, default: None
822
- The root directory under which `current.model.loaded` will store loaded models
823
-
824
-
825
-
826
- """
827
- ...
828
-
829
- @typing.overload
830
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
831
- """
832
- Decorator prototype for all step decorators. This function gets specialized
833
- and imported for all decorators types by _import_plugin_decorators().
834
- """
835
- ...
836
-
837
- @typing.overload
838
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
839
- ...
840
-
841
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
842
- """
843
- Decorator prototype for all step decorators. This function gets specialized
844
- and imported for all decorators types by _import_plugin_decorators().
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.
845
830
  """
846
831
  ...
847
832
 
@@ -879,251 +864,101 @@ def secrets(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None],
879
864
  ...
880
865
 
881
866
  @typing.overload
882
- 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]]]:
883
- """
884
- Specifies the Conda environment for the step.
885
-
886
- Information in this decorator will augment any
887
- attributes set in the `@conda_base` flow-level decorator. Hence,
888
- you can use `@conda_base` to set packages required by all
889
- steps and use `@conda` to specify step-specific overrides.
890
-
891
- Parameters
892
- ----------
893
- packages : Dict[str, str], default {}
894
- Packages to use for this step. The key is the name of the package
895
- and the value is the version to use.
896
- libraries : Dict[str, str], default {}
897
- Supported for backward compatibility. When used with packages, packages will take precedence.
898
- python : str, optional, default None
899
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
900
- that the version used will correspond to the version of the Python interpreter used to start the run.
901
- disabled : bool, default False
902
- If set to True, disables @conda.
903
- """
904
- ...
905
-
906
- @typing.overload
907
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
908
- ...
909
-
910
- @typing.overload
911
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
912
- ...
913
-
914
- 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):
915
- """
916
- Specifies the Conda environment for the step.
917
-
918
- Information in this decorator will augment any
919
- attributes set in the `@conda_base` flow-level decorator. Hence,
920
- you can use `@conda_base` to set packages required by all
921
- steps and use `@conda` to specify step-specific overrides.
922
-
923
- Parameters
924
- ----------
925
- packages : Dict[str, str], default {}
926
- Packages to use for this step. The key is the name of the package
927
- and the value is the version to use.
928
- libraries : Dict[str, str], default {}
929
- Supported for backward compatibility. When used with packages, packages will take precedence.
930
- python : str, optional, default None
931
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
932
- that the version used will correspond to the version of the Python interpreter used to start the run.
933
- disabled : bool, default False
934
- If set to True, disables @conda.
935
- """
936
- ...
937
-
938
- @typing.overload
939
- 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]]]:
940
- """
941
- Specifies a timeout for your step.
942
-
943
- This decorator is useful if this step may hang indefinitely.
944
-
945
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
946
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
947
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
948
-
949
- Note that all the values specified in parameters are added together so if you specify
950
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
951
-
952
- Parameters
953
- ----------
954
- seconds : int, default 0
955
- Number of seconds to wait prior to timing out.
956
- minutes : int, default 0
957
- Number of minutes to wait prior to timing out.
958
- hours : int, default 0
959
- Number of hours to wait prior to timing out.
960
- """
961
- ...
962
-
963
- @typing.overload
964
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
965
- ...
966
-
967
- @typing.overload
968
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
969
- ...
970
-
971
- 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):
972
- """
973
- Specifies a timeout for your step.
974
-
975
- This decorator is useful if this step may hang indefinitely.
976
-
977
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
978
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
979
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
980
-
981
- Note that all the values specified in parameters are added together so if you specify
982
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
983
-
984
- Parameters
985
- ----------
986
- seconds : int, default 0
987
- Number of seconds to wait prior to timing out.
988
- minutes : int, default 0
989
- Number of minutes to wait prior to timing out.
990
- hours : int, default 0
991
- Number of hours to wait prior to timing out.
992
- """
993
- ...
994
-
995
- @typing.overload
996
- 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]]]:
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]]]:
997
868
  """
998
- Specifies the number of times the task corresponding
999
- to a step needs to be retried.
1000
-
1001
- This decorator is useful for handling transient errors, such as networking issues.
1002
- If your task contains operations that can't be retried safely, e.g. database updates,
1003
- it is advisable to annotate it with `@retry(times=0)`.
1004
-
1005
- This can be used in conjunction with the `@catch` decorator. The `@catch`
1006
- decorator will execute a no-op task after all retries have been exhausted,
1007
- ensuring that the flow execution can continue.
869
+ Specifies environment variables to be set prior to the execution of a step.
1008
870
 
1009
871
  Parameters
1010
872
  ----------
1011
- times : int, default 3
1012
- Number of times to retry this task.
1013
- minutes_between_retries : int, default 2
1014
- Number of minutes between retries.
873
+ vars : Dict[str, str], default {}
874
+ Dictionary of environment variables to set.
1015
875
  """
1016
876
  ...
1017
877
 
1018
878
  @typing.overload
1019
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
879
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1020
880
  ...
1021
881
 
1022
882
  @typing.overload
1023
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
883
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1024
884
  ...
1025
885
 
1026
- 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):
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] = {}):
1027
887
  """
1028
- Specifies the number of times the task corresponding
1029
- to a step needs to be retried.
1030
-
1031
- This decorator is useful for handling transient errors, such as networking issues.
1032
- If your task contains operations that can't be retried safely, e.g. database updates,
1033
- it is advisable to annotate it with `@retry(times=0)`.
1034
-
1035
- This can be used in conjunction with the `@catch` decorator. The `@catch`
1036
- decorator will execute a no-op task after all retries have been exhausted,
1037
- ensuring that the flow execution can continue.
888
+ Specifies environment variables to be set prior to the execution of a step.
1038
889
 
1039
890
  Parameters
1040
891
  ----------
1041
- times : int, default 3
1042
- Number of times to retry this task.
1043
- minutes_between_retries : int, default 2
1044
- Number of minutes between retries.
892
+ vars : Dict[str, str], default {}
893
+ Dictionary of environment variables to set.
1045
894
  """
1046
895
  ...
1047
896
 
1048
897
  @typing.overload
1049
- 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]]]:
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]]]:
1050
899
  """
1051
- Creates a human-readable report, a Metaflow Card, after this step completes.
900
+ Specifies the PyPI packages for the step.
1052
901
 
1053
- Note that you may add multiple `@card` decorators in a step with different parameters.
902
+ Information in this decorator will augment any
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.
1054
906
 
1055
907
  Parameters
1056
908
  ----------
1057
- type : str, default 'default'
1058
- Card type.
1059
- id : str, optional, default None
1060
- If multiple cards are present, use this id to identify this card.
1061
- options : Dict[str, Any], default {}
1062
- Options passed to the card. The contents depend on the card type.
1063
- timeout : int, default 45
1064
- Interrupt reporting if it takes more than this many seconds.
1065
-
1066
-
909
+ packages : Dict[str, str], default: {}
910
+ Packages to use for this step. The key is the name of the package
911
+ and the value is the version to use.
912
+ python : str, optional, default: None
913
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
914
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1067
915
  """
1068
916
  ...
1069
917
 
1070
918
  @typing.overload
1071
- def card(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]:
1072
920
  ...
1073
921
 
1074
922
  @typing.overload
1075
- def card(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]:
1076
924
  ...
1077
925
 
1078
- 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):
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):
1079
927
  """
1080
- Creates a human-readable report, a Metaflow Card, after this step completes.
928
+ Specifies the PyPI packages for the step.
1081
929
 
1082
- Note that you may add multiple `@card` decorators in a step with different parameters.
930
+ Information in this decorator will augment any
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.
1083
934
 
1084
935
  Parameters
1085
936
  ----------
1086
- type : str, default 'default'
1087
- Card type.
1088
- id : str, optional, default None
1089
- If multiple cards are present, use this id to identify this card.
1090
- options : Dict[str, Any], default {}
1091
- Options passed to the card. The contents depend on the card type.
1092
- timeout : int, default 45
1093
- Interrupt reporting if it takes more than this many seconds.
1094
-
1095
-
937
+ packages : Dict[str, str], default: {}
938
+ Packages to use for this step. The key is the name of the package
939
+ and the value is the version to use.
940
+ python : str, optional, default: None
941
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
942
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1096
943
  """
1097
944
  ...
1098
945
 
1099
946
  @typing.overload
1100
- 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]:
1101
948
  """
1102
- Specifies environment variables to be set prior to the execution of a step.
1103
-
1104
- Parameters
1105
- ----------
1106
- vars : Dict[str, str], default {}
1107
- 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().
1108
951
  """
1109
952
  ...
1110
953
 
1111
954
  @typing.overload
1112
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1113
- ...
1114
-
1115
- @typing.overload
1116
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
955
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1117
956
  ...
1118
957
 
1119
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
958
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1120
959
  """
1121
- Specifies environment variables to be set prior to the execution of a step.
1122
-
1123
- Parameters
1124
- ----------
1125
- vars : Dict[str, str], default {}
1126
- Dictionary of environment variables to set.
960
+ Decorator prototype for all step decorators. This function gets specialized
961
+ and imported for all decorators types by _import_plugin_decorators().
1127
962
  """
1128
963
  ...
1129
964
 
@@ -1194,200 +1029,225 @@ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: ty
1194
1029
  ...
1195
1030
 
1196
1031
  @typing.overload
1197
- def schedule(*, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1032
+ def checkpoint(*, load_policy: str = "fresh", 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]]]:
1198
1033
  """
1199
- Specifies the times when the flow should be run when running on a
1200
- production scheduler.
1034
+ Enables checkpointing for a step.
1035
+
1201
1036
 
1202
1037
  Parameters
1203
1038
  ----------
1204
- hourly : bool, default False
1205
- Run the workflow hourly.
1206
- daily : bool, default True
1207
- Run the workflow daily.
1208
- weekly : bool, default False
1209
- Run the workflow weekly.
1210
- cron : str, optional, default None
1211
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1212
- specified by this expression.
1213
- timezone : str, optional, default None
1214
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1215
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1216
- """
1217
- ...
1218
-
1039
+ load_policy : str, default: "fresh"
1040
+ The policy for loading the checkpoint. The following policies are supported:
1041
+ - "eager": Loads the the latest available checkpoint within the namespace.
1042
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
1043
+ will be loaded at the start of the task.
1044
+ - "none": Do not load any checkpoint
1045
+ - "fresh": Loads the lastest checkpoint created within the running Task.
1046
+ This mode helps loading checkpoints across various retry attempts of the same task.
1047
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
1048
+ created within the task will be loaded when the task is retries execution on failure.
1049
+
1050
+ temp_dir_root : str, default: None
1051
+ The root directory under which `current.checkpoint.directory` will be created.
1052
+
1053
+
1054
+
1055
+ """
1056
+ ...
1057
+
1219
1058
  @typing.overload
1220
- def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1059
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1221
1060
  ...
1222
1061
 
1223
- def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None):
1062
+ @typing.overload
1063
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1064
+ ...
1065
+
1066
+ def checkpoint(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, load_policy: str = "fresh", temp_dir_root: str = None):
1224
1067
  """
1225
- Specifies the times when the flow should be run when running on a
1226
- production scheduler.
1068
+ Enables checkpointing for a step.
1069
+
1227
1070
 
1228
1071
  Parameters
1229
1072
  ----------
1230
- hourly : bool, default False
1231
- Run the workflow hourly.
1232
- daily : bool, default True
1233
- Run the workflow daily.
1234
- weekly : bool, default False
1235
- Run the workflow weekly.
1236
- cron : str, optional, default None
1237
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1238
- specified by this expression.
1239
- timezone : str, optional, default None
1240
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1241
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1073
+ load_policy : str, default: "fresh"
1074
+ The policy for loading the checkpoint. The following policies are supported:
1075
+ - "eager": Loads the the latest available checkpoint within the namespace.
1076
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
1077
+ will be loaded at the start of the task.
1078
+ - "none": Do not load any checkpoint
1079
+ - "fresh": Loads the lastest checkpoint created within the running Task.
1080
+ This mode helps loading checkpoints across various retry attempts of the same task.
1081
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
1082
+ created within the task will be loaded when the task is retries execution on failure.
1083
+
1084
+ temp_dir_root : str, default: None
1085
+ The root directory under which `current.checkpoint.directory` will be created.
1086
+
1087
+
1088
+
1242
1089
  """
1243
1090
  ...
1244
1091
 
1245
1092
  @typing.overload
1246
- 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]]:
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]]]:
1247
1094
  """
1248
- Specifies the event(s) that this flow depends on.
1249
-
1250
- ```
1251
- @trigger(event='foo')
1252
- ```
1253
- or
1254
- ```
1255
- @trigger(events=['foo', 'bar'])
1256
- ```
1257
-
1258
- Additionally, you can specify the parameter mappings
1259
- to map event payload to Metaflow parameters for the flow.
1260
- ```
1261
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1262
- ```
1263
- or
1264
- ```
1265
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1266
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1267
- ```
1095
+ Creates a human-readable report, a Metaflow Card, after this step completes.
1268
1096
 
1269
- 'parameters' can also be a list of strings and tuples like so:
1270
- ```
1271
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1272
- ```
1273
- This is equivalent to:
1274
- ```
1275
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1276
- ```
1097
+ Note that you may add multiple `@card` decorators in a step with different parameters.
1277
1098
 
1278
1099
  Parameters
1279
1100
  ----------
1280
- event : Union[str, Dict[str, Any]], optional, default None
1281
- Event dependency for this flow.
1282
- events : List[Union[str, Dict[str, Any]]], default []
1283
- Events dependency for this flow.
1101
+ type : str, default 'default'
1102
+ Card type.
1103
+ id : str, optional, default None
1104
+ If multiple cards are present, use this id to identify this card.
1284
1105
  options : Dict[str, Any], default {}
1285
- Backend-specific configuration for tuning eventing behavior.
1106
+ Options passed to the card. The contents depend on the card type.
1107
+ timeout : int, default 45
1108
+ Interrupt reporting if it takes more than this many seconds.
1286
1109
 
1287
1110
 
1288
1111
  """
1289
1112
  ...
1290
1113
 
1291
1114
  @typing.overload
1292
- def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1115
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1293
1116
  ...
1294
1117
 
1295
- 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] = {}):
1118
+ @typing.overload
1119
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1120
+ ...
1121
+
1122
+ 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):
1296
1123
  """
1297
- Specifies the event(s) that this flow depends on.
1298
-
1299
- ```
1300
- @trigger(event='foo')
1301
- ```
1302
- or
1303
- ```
1304
- @trigger(events=['foo', 'bar'])
1305
- ```
1306
-
1307
- Additionally, you can specify the parameter mappings
1308
- to map event payload to Metaflow parameters for the flow.
1309
- ```
1310
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1311
- ```
1312
- or
1313
- ```
1314
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1315
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1316
- ```
1124
+ Creates a human-readable report, a Metaflow Card, after this step completes.
1317
1125
 
1318
- 'parameters' can also be a list of strings and tuples like so:
1319
- ```
1320
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1321
- ```
1322
- This is equivalent to:
1323
- ```
1324
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1325
- ```
1126
+ Note that you may add multiple `@card` decorators in a step with different parameters.
1326
1127
 
1327
1128
  Parameters
1328
1129
  ----------
1329
- event : Union[str, Dict[str, Any]], optional, default None
1330
- Event dependency for this flow.
1331
- events : List[Union[str, Dict[str, Any]]], default []
1332
- Events dependency for this flow.
1130
+ type : str, default 'default'
1131
+ Card type.
1132
+ id : str, optional, default None
1133
+ If multiple cards are present, use this id to identify this card.
1333
1134
  options : Dict[str, Any], default {}
1334
- Backend-specific configuration for tuning eventing behavior.
1135
+ Options passed to the card. The contents depend on the card type.
1136
+ timeout : int, default 45
1137
+ Interrupt reporting if it takes more than this many seconds.
1335
1138
 
1336
1139
 
1337
1140
  """
1338
1141
  ...
1339
1142
 
1340
- def nim(*, models: "list[NIM]", backend: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1143
+ @typing.overload
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]]]:
1341
1145
  """
1342
- This decorator is used to run NIM containers in Metaflow tasks as sidecars.
1146
+ Specifies the number of times the task corresponding
1147
+ to a step needs to be retried.
1343
1148
 
1344
- User code call
1345
- -----------
1346
- @nim(
1347
- models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
1348
- backend='managed'
1349
- )
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)`.
1350
1152
 
1351
- Valid backend options
1352
- ---------------------
1353
- - 'managed': Outerbounds selects a compute provider based on the model.
1354
- - 🚧 'dataplane': Run in your account.
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.
1355
1156
 
1356
- Valid model options
1357
- ----------------
1358
- - 'meta/llama3-8b-instruct': 8B parameter model
1359
- - 'meta/llama3-70b-instruct': 70B parameter model
1360
- - Upon request, any model here: https://nvcf.ngc.nvidia.com/functions?filter=nvidia-functions
1157
+ Parameters
1158
+ ----------
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.
1163
+ """
1164
+ ...
1165
+
1166
+ @typing.overload
1167
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1168
+ ...
1169
+
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):
1175
+ """
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.
1361
1186
 
1362
1187
  Parameters
1363
1188
  ----------
1364
- models: list[NIM]
1365
- List of NIM containers running models in sidecars.
1366
- backend: str
1367
- Compute provider to run the NIM container.
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.
1368
1193
  """
1369
1194
  ...
1370
1195
 
1371
- def project(*, name: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1196
+ @typing.overload
1197
+ def trigger_on_finish(*, flow: typing.Union[typing.Dict[str, 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]]:
1372
1198
  """
1373
- Specifies what flows belong to the same project.
1199
+ Specifies the flow(s) that this flow depends on.
1374
1200
 
1375
- A project-specific namespace is created for all flows that
1376
- use the same `@project(name)`.
1201
+ ```
1202
+ @trigger_on_finish(flow='FooFlow')
1203
+ ```
1204
+ or
1205
+ ```
1206
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1207
+ ```
1208
+ This decorator respects the @project decorator and triggers the flow
1209
+ when upstream runs within the same namespace complete successfully
1210
+
1211
+ Additionally, you can specify project aware upstream flow dependencies
1212
+ by specifying the fully qualified project_flow_name.
1213
+ ```
1214
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1215
+ ```
1216
+ or
1217
+ ```
1218
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1219
+ ```
1220
+
1221
+ You can also specify just the project or project branch (other values will be
1222
+ inferred from the current project or project branch):
1223
+ ```
1224
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1225
+ ```
1226
+
1227
+ Note that `branch` is typically one of:
1228
+ - `prod`
1229
+ - `user.bob`
1230
+ - `test.my_experiment`
1231
+ - `prod.staging`
1377
1232
 
1378
1233
  Parameters
1379
1234
  ----------
1380
- name : str
1381
- Project name. Make sure that the name is unique amongst all
1382
- projects that use the same production scheduler. The name may
1383
- contain only lowercase alphanumeric characters and underscores.
1235
+ flow : Union[str, Dict[str, str]], optional, default None
1236
+ Upstream flow dependency for this flow.
1237
+ flows : List[Union[str, Dict[str, str]]], default []
1238
+ Upstream flow dependencies for this flow.
1239
+ options : Dict[str, Any], default {}
1240
+ Backend-specific configuration for tuning eventing behavior.
1384
1241
 
1385
1242
 
1386
1243
  """
1387
1244
  ...
1388
1245
 
1389
1246
  @typing.overload
1390
- def trigger_on_finish(*, flow: typing.Union[typing.Dict[str, 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]]:
1247
+ def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1248
+ ...
1249
+
1250
+ def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}):
1391
1251
  """
1392
1252
  Specifies the flow(s) that this flow depends on.
1393
1253
 
@@ -1417,75 +1277,162 @@ def trigger_on_finish(*, flow: typing.Union[typing.Dict[str, str], str, None] =
1417
1277
  @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1418
1278
  ```
1419
1279
 
1420
- Note that `branch` is typically one of:
1421
- - `prod`
1422
- - `user.bob`
1423
- - `test.my_experiment`
1424
- - `prod.staging`
1280
+ Note that `branch` is typically one of:
1281
+ - `prod`
1282
+ - `user.bob`
1283
+ - `test.my_experiment`
1284
+ - `prod.staging`
1285
+
1286
+ Parameters
1287
+ ----------
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.
1292
+ options : Dict[str, Any], default {}
1293
+ Backend-specific configuration for tuning eventing behavior.
1294
+
1295
+
1296
+ """
1297
+ ...
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
+
1341
+ def nim(*, models: "list[NIM]", backend: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1342
+ """
1343
+ This decorator is used to run NIM containers in Metaflow tasks as sidecars.
1344
+
1345
+ User code call
1346
+ -----------
1347
+ @nim(
1348
+ models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
1349
+ backend='managed'
1350
+ )
1351
+
1352
+ Valid backend options
1353
+ ---------------------
1354
+ - 'managed': Outerbounds selects a compute provider based on the model.
1355
+ - 🚧 'dataplane': Run in your account.
1356
+
1357
+ Valid model options
1358
+ ----------------
1359
+ - 'meta/llama3-8b-instruct': 8B parameter model
1360
+ - 'meta/llama3-70b-instruct': 70B parameter model
1361
+ - Upon request, any model here: https://nvcf.ngc.nvidia.com/functions?filter=nvidia-functions
1362
+
1363
+ Parameters
1364
+ ----------
1365
+ models: list[NIM]
1366
+ List of NIM containers running models in sidecars.
1367
+ backend: str
1368
+ Compute provider to run the NIM container.
1369
+ """
1370
+ ...
1371
+
1372
+ def project(*, name: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1373
+ """
1374
+ Specifies what flows belong to the same project.
1375
+
1376
+ A project-specific namespace is created for all flows that
1377
+ use the same `@project(name)`.
1425
1378
 
1426
1379
  Parameters
1427
1380
  ----------
1428
- flow : Union[str, Dict[str, str]], optional, default None
1429
- Upstream flow dependency for this flow.
1430
- flows : List[Union[str, Dict[str, str]]], default []
1431
- Upstream flow dependencies for this flow.
1432
- options : Dict[str, Any], default {}
1433
- Backend-specific configuration for tuning eventing behavior.
1381
+ name : str
1382
+ Project name. Make sure that the name is unique amongst all
1383
+ projects that use the same production scheduler. The name may
1384
+ contain only lowercase alphanumeric characters and underscores.
1434
1385
 
1435
1386
 
1436
1387
  """
1437
1388
  ...
1438
1389
 
1439
1390
  @typing.overload
1440
- def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1391
+ def schedule(*, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1392
+ """
1393
+ Specifies the times when the flow should be run when running on a
1394
+ production scheduler.
1395
+
1396
+ Parameters
1397
+ ----------
1398
+ hourly : bool, default False
1399
+ Run the workflow hourly.
1400
+ daily : bool, default True
1401
+ Run the workflow daily.
1402
+ weekly : bool, default False
1403
+ Run the workflow weekly.
1404
+ cron : str, optional, default None
1405
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1406
+ specified by this expression.
1407
+ timezone : str, optional, default None
1408
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1409
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1410
+ """
1441
1411
  ...
1442
1412
 
1443
- def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}):
1413
+ @typing.overload
1414
+ def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1415
+ ...
1416
+
1417
+ def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None):
1444
1418
  """
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`
1419
+ Specifies the times when the flow should be run when running on a
1420
+ production scheduler.
1478
1421
 
1479
1422
  Parameters
1480
1423
  ----------
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
-
1424
+ hourly : bool, default False
1425
+ Run the workflow hourly.
1426
+ daily : bool, default True
1427
+ Run the workflow daily.
1428
+ weekly : bool, default False
1429
+ Run the workflow weekly.
1430
+ cron : str, optional, default None
1431
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1432
+ specified by this expression.
1433
+ timezone : str, optional, default None
1434
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1435
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1489
1436
  """
1490
1437
  ...
1491
1438
 
@@ -1528,48 +1475,6 @@ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packag
1528
1475
  """
1529
1476
  ...
1530
1477
 
1531
- 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]]:
1532
- """
1533
- 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)
1534
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1535
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1536
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
1537
- starts only after all sensors finish.
1538
-
1539
- Parameters
1540
- ----------
1541
- timeout : int
1542
- Time, in seconds before the task times out and fails. (Default: 3600)
1543
- poke_interval : int
1544
- Time in seconds that the job should wait in between each try. (Default: 60)
1545
- mode : str
1546
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1547
- exponential_backoff : bool
1548
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1549
- pool : str
1550
- the slot pool this task should run in,
1551
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1552
- soft_fail : bool
1553
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1554
- name : str
1555
- Name of the sensor on Airflow
1556
- description : str
1557
- Description of sensor in the Airflow UI
1558
- bucket_key : Union[str, List[str]]
1559
- The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1560
- When it's specified as a full s3:// url, please leave `bucket_name` as None
1561
- bucket_name : str
1562
- Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1563
- When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1564
- wildcard_match : bool
1565
- whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1566
- aws_conn_id : str
1567
- a reference to the s3 connection on Airflow. (Default: None)
1568
- verify : bool
1569
- Whether or not to verify SSL certificates for S3 connection. (Default: None)
1570
- """
1571
- ...
1572
-
1573
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]]:
1574
1479
  """
1575
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.
@@ -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