ob-metaflow-stubs 6.0.3.115__py2.py3-none-any.whl → 6.0.3.116__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 +587 -587
  2. metaflow-stubs/cards.pyi +5 -5
  3. metaflow-stubs/cli.pyi +2 -2
  4. metaflow-stubs/client/__init__.pyi +2 -2
  5. metaflow-stubs/client/core.pyi +4 -4
  6. metaflow-stubs/client/filecache.pyi +1 -1
  7. metaflow-stubs/clone_util.pyi +1 -1
  8. metaflow-stubs/events.pyi +1 -1
  9. metaflow-stubs/exception.pyi +1 -1
  10. metaflow-stubs/flowspec.pyi +4 -4
  11. metaflow-stubs/generated_for.txt +1 -1
  12. metaflow-stubs/includefile.pyi +3 -3
  13. metaflow-stubs/info_file.pyi +1 -1
  14. metaflow-stubs/metadata/metadata.pyi +1 -1
  15. metaflow-stubs/metadata/util.pyi +1 -1
  16. metaflow-stubs/metaflow_config.pyi +1 -1
  17. metaflow-stubs/metaflow_current.pyi +76 -76
  18. metaflow-stubs/mflog/mflog.pyi +1 -1
  19. metaflow-stubs/multicore_utils.pyi +1 -1
  20. metaflow-stubs/parameters.pyi +2 -2
  21. metaflow-stubs/plugins/__init__.pyi +1 -1
  22. metaflow-stubs/plugins/airflow/__init__.pyi +1 -1
  23. metaflow-stubs/plugins/airflow/airflow_utils.pyi +1 -1
  24. metaflow-stubs/plugins/airflow/exception.pyi +1 -1
  25. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +1 -1
  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 +1 -1
  30. metaflow-stubs/plugins/argo/argo_client.pyi +1 -1
  31. metaflow-stubs/plugins/argo/argo_events.pyi +1 -1
  32. metaflow-stubs/plugins/argo/argo_workflows.pyi +4 -4
  33. metaflow-stubs/plugins/argo/argo_workflows_cli.pyi +5 -5
  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 +1 -1
  37. metaflow-stubs/plugins/aws/aws_client.pyi +1 -1
  38. metaflow-stubs/plugins/aws/aws_utils.pyi +1 -1
  39. metaflow-stubs/plugins/aws/batch/__init__.pyi +1 -1
  40. metaflow-stubs/plugins/aws/batch/batch.pyi +1 -1
  41. metaflow-stubs/plugins/aws/batch/batch_client.pyi +1 -1
  42. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +2 -2
  43. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +1 -1
  44. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +2 -2
  45. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +1 -1
  46. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +1 -1
  47. metaflow-stubs/plugins/aws/step_functions/production_token.pyi +1 -1
  48. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +1 -1
  49. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +1 -1
  50. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +1 -1
  51. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer.pyi +1 -1
  52. metaflow-stubs/plugins/azure/__init__.pyi +1 -1
  53. metaflow-stubs/plugins/azure/azure_credential.pyi +1 -1
  54. metaflow-stubs/plugins/azure/azure_exceptions.pyi +1 -1
  55. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +2 -2
  56. metaflow-stubs/plugins/azure/azure_utils.pyi +1 -1
  57. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +1 -1
  58. metaflow-stubs/plugins/azure/includefile_support.pyi +1 -1
  59. metaflow-stubs/plugins/cards/__init__.pyi +1 -1
  60. metaflow-stubs/plugins/cards/card_cli.pyi +3 -3
  61. metaflow-stubs/plugins/cards/card_client.pyi +2 -2
  62. metaflow-stubs/plugins/cards/card_creator.pyi +1 -1
  63. metaflow-stubs/plugins/cards/card_datastore.pyi +1 -1
  64. metaflow-stubs/plugins/cards/card_decorator.pyi +2 -2
  65. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +1 -1
  66. metaflow-stubs/plugins/cards/card_modules/basic.pyi +2 -2
  67. metaflow-stubs/plugins/cards/card_modules/card.pyi +1 -1
  68. metaflow-stubs/plugins/cards/card_modules/components.pyi +3 -3
  69. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +1 -1
  70. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +1 -1
  71. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +1 -1
  72. metaflow-stubs/plugins/cards/card_resolver.pyi +1 -1
  73. metaflow-stubs/plugins/cards/component_serializer.pyi +2 -2
  74. metaflow-stubs/plugins/cards/exception.pyi +1 -1
  75. metaflow-stubs/plugins/catch_decorator.pyi +2 -2
  76. metaflow-stubs/plugins/datatools/__init__.pyi +2 -2
  77. metaflow-stubs/plugins/datatools/local.pyi +1 -1
  78. metaflow-stubs/plugins/datatools/s3/__init__.pyi +2 -2
  79. metaflow-stubs/plugins/datatools/s3/s3.pyi +5 -5
  80. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +1 -1
  81. metaflow-stubs/plugins/datatools/s3/s3util.pyi +1 -1
  82. metaflow-stubs/plugins/debug_logger.pyi +1 -1
  83. metaflow-stubs/plugins/debug_monitor.pyi +1 -1
  84. metaflow-stubs/plugins/environment_decorator.pyi +1 -1
  85. metaflow-stubs/plugins/events_decorator.pyi +2 -2
  86. metaflow-stubs/plugins/frameworks/__init__.pyi +1 -1
  87. metaflow-stubs/plugins/frameworks/pytorch.pyi +2 -2
  88. metaflow-stubs/plugins/gcp/__init__.pyi +1 -1
  89. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +2 -2
  90. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +1 -1
  91. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +1 -1
  92. metaflow-stubs/plugins/gcp/gs_utils.pyi +1 -1
  93. metaflow-stubs/plugins/gcp/includefile_support.pyi +1 -1
  94. metaflow-stubs/plugins/kubernetes/__init__.pyi +1 -1
  95. metaflow-stubs/plugins/kubernetes/kube_utils.pyi +1 -1
  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 +1 -1
  99. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +2 -2
  100. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +1 -1
  101. metaflow-stubs/plugins/logs_cli.pyi +3 -3
  102. metaflow-stubs/plugins/package_cli.pyi +1 -1
  103. metaflow-stubs/plugins/parallel_decorator.pyi +2 -2
  104. metaflow-stubs/plugins/perimeters.pyi +1 -1
  105. metaflow-stubs/plugins/project_decorator.pyi +2 -2
  106. metaflow-stubs/plugins/pypi/__init__.pyi +1 -1
  107. metaflow-stubs/plugins/pypi/conda_decorator.pyi +1 -1
  108. metaflow-stubs/plugins/pypi/conda_environment.pyi +3 -3
  109. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +1 -1
  110. metaflow-stubs/plugins/pypi/pypi_environment.pyi +2 -2
  111. metaflow-stubs/plugins/pypi/utils.pyi +1 -1
  112. metaflow-stubs/plugins/resources_decorator.pyi +1 -1
  113. metaflow-stubs/plugins/retry_decorator.pyi +1 -1
  114. metaflow-stubs/plugins/secrets/__init__.pyi +1 -1
  115. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +1 -1
  116. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +1 -1
  117. metaflow-stubs/plugins/storage_executor.pyi +1 -1
  118. metaflow-stubs/plugins/tag_cli.pyi +3 -3
  119. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +2 -2
  120. metaflow-stubs/plugins/timeout_decorator.pyi +2 -2
  121. metaflow-stubs/procpoll.pyi +1 -1
  122. metaflow-stubs/profilers/__init__.pyi +1 -1
  123. metaflow-stubs/pylint_wrapper.pyi +1 -1
  124. metaflow-stubs/runner/__init__.pyi +1 -1
  125. metaflow-stubs/runner/deployer.pyi +2 -2
  126. metaflow-stubs/runner/metaflow_runner.pyi +4 -4
  127. metaflow-stubs/runner/nbdeploy.pyi +1 -1
  128. metaflow-stubs/runner/nbrun.pyi +1 -1
  129. metaflow-stubs/runner/subprocess_manager.pyi +1 -1
  130. metaflow-stubs/runner/utils.pyi +1 -1
  131. metaflow-stubs/system/__init__.pyi +2 -2
  132. metaflow-stubs/system/system_logger.pyi +1 -1
  133. metaflow-stubs/system/system_monitor.pyi +2 -2
  134. metaflow-stubs/tagging_util.pyi +1 -1
  135. metaflow-stubs/tuple_util.pyi +1 -1
  136. {ob_metaflow_stubs-6.0.3.115.dist-info → ob_metaflow_stubs-6.0.3.116.dist-info}/METADATA +1 -1
  137. ob_metaflow_stubs-6.0.3.116.dist-info/RECORD +140 -0
  138. ob_metaflow_stubs-6.0.3.115.dist-info/RECORD +0 -140
  139. {ob_metaflow_stubs-6.0.3.115.dist-info → ob_metaflow_stubs-6.0.3.116.dist-info}/WHEEL +0 -0
  140. {ob_metaflow_stubs-6.0.3.115.dist-info → ob_metaflow_stubs-6.0.3.116.dist-info}/top_level.txt +0 -0
@@ -1,24 +1,24 @@
1
1
  ##################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
3
  # MF version: 2.12.27.1+obcheckpoint(0.1.1);ob(v1) #
4
- # Generated on 2024-10-25T21:39:54.730102 #
4
+ # Generated on 2024-11-04T20:32:12.972782 #
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.datastore.inputs
11
12
  import typing
13
+ import metaflow.flowspec
12
14
  import metaflow.metaflow_current
13
- import datetime
14
- import metaflow._vendor.click.types
15
15
  import metaflow.runner.metaflow_runner
16
- import metaflow.parameters
17
- import metaflow.events
18
- import metaflow_extensions.obcheckpoint.plugins.machine_learning_utilities.datastructures
16
+ import datetime
19
17
  import metaflow.client.core
20
- import metaflow.datastore.inputs
21
- import metaflow.flowspec
18
+ import metaflow_extensions.obcheckpoint.plugins.machine_learning_utilities.datastructures
19
+ import metaflow.events
20
+ import metaflow.parameters
21
+ import metaflow._vendor.click.types
22
22
  FlowSpecDerived = typing.TypeVar("FlowSpecDerived", bound="FlowSpec", contravariant=False, covariant=False)
23
23
  StepFlag = typing.NewType("StepFlag", bool)
24
24
 
@@ -513,174 +513,55 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
513
513
  ...
514
514
 
515
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
-
532
- @typing.overload
533
- def secrets(*, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = []) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
534
- """
535
- Specifies secrets to be retrieved and injected as environment variables prior to
536
- the execution of a step.
537
-
538
- Parameters
539
- ----------
540
- sources : List[Union[str, Dict[str, Any]]], default: []
541
- List of secret specs, defining how the secrets are to be retrieved
542
- """
543
- ...
544
-
545
- @typing.overload
546
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
547
- ...
548
-
549
- @typing.overload
550
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
551
- ...
552
-
553
- def secrets(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = []):
516
+ 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]]]:
554
517
  """
555
- Specifies secrets to be retrieved and injected as environment variables prior to
556
- the execution of a step.
518
+ Specifies the number of times the task corresponding
519
+ to a step needs to be retried.
557
520
 
558
- Parameters
559
- ----------
560
- sources : List[Union[str, Dict[str, Any]]], default: []
561
- List of secret specs, defining how the secrets are to be retrieved
562
- """
563
- ...
564
-
565
- @typing.overload
566
- 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]]]:
567
- """
568
- Enables checkpointing for a step.
521
+ This decorator is useful for handling transient errors, such as networking issues.
522
+ If your task contains operations that can't be retried safely, e.g. database updates,
523
+ it is advisable to annotate it with `@retry(times=0)`.
569
524
 
525
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
526
+ decorator will execute a no-op task after all retries have been exhausted,
527
+ ensuring that the flow execution can continue.
570
528
 
571
529
  Parameters
572
530
  ----------
573
- load_policy : str, default: "fresh"
574
- The policy for loading the checkpoint. The following policies are supported:
575
- - "eager": Loads the the latest available checkpoint within the namespace.
576
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
577
- will be loaded at the start of the task.
578
- - "none": Do not load any checkpoint
579
- - "fresh": Loads the lastest checkpoint created within the running Task.
580
- This mode helps loading checkpoints across various retry attempts of the same task.
581
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
582
- created within the task will be loaded when the task is retries execution on failure.
583
-
584
- temp_dir_root : str, default: None
585
- The root directory under which `current.checkpoint.directory` will be created.
586
-
587
-
588
-
531
+ times : int, default 3
532
+ Number of times to retry this task.
533
+ minutes_between_retries : int, default 2
534
+ Number of minutes between retries.
589
535
  """
590
536
  ...
591
537
 
592
538
  @typing.overload
593
- def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
539
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
594
540
  ...
595
541
 
596
542
  @typing.overload
597
- def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
598
- ...
599
-
600
- 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):
601
- """
602
- Enables checkpointing for a step.
603
-
604
-
605
- Parameters
606
- ----------
607
- load_policy : str, default: "fresh"
608
- The policy for loading the checkpoint. The following policies are supported:
609
- - "eager": Loads the the latest available checkpoint within the namespace.
610
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
611
- will be loaded at the start of the task.
612
- - "none": Do not load any checkpoint
613
- - "fresh": Loads the lastest checkpoint created within the running Task.
614
- This mode helps loading checkpoints across various retry attempts of the same task.
615
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
616
- created within the task will be loaded when the task is retries execution on failure.
617
-
618
- temp_dir_root : str, default: None
619
- The root directory under which `current.checkpoint.directory` will be created.
620
-
621
-
622
-
623
- """
543
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
624
544
  ...
625
545
 
626
- @typing.overload
627
- 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]]]:
546
+ 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):
628
547
  """
629
- Enables loading / saving of models within a step.
630
-
631
-
632
- Parameters
633
- ----------
634
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
635
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
636
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
637
- - `current.checkpoint`
638
- - `current.model`
639
- - `current.huggingface_hub`
640
-
641
- 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
642
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
643
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
644
-
645
- temp_dir_root : str, default: None
646
- The root directory under which `current.model.loaded` will store loaded models
647
-
648
-
548
+ Specifies the number of times the task corresponding
549
+ to a step needs to be retried.
649
550
 
650
- """
651
- ...
652
-
653
- @typing.overload
654
- def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
655
- ...
656
-
657
- @typing.overload
658
- def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
659
- ...
660
-
661
- 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):
662
- """
663
- Enables loading / saving of models within a step.
551
+ This decorator is useful for handling transient errors, such as networking issues.
552
+ If your task contains operations that can't be retried safely, e.g. database updates,
553
+ it is advisable to annotate it with `@retry(times=0)`.
664
554
 
555
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
556
+ decorator will execute a no-op task after all retries have been exhausted,
557
+ ensuring that the flow execution can continue.
665
558
 
666
559
  Parameters
667
560
  ----------
668
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
669
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
670
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
671
- - `current.checkpoint`
672
- - `current.model`
673
- - `current.huggingface_hub`
674
-
675
- 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
676
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
677
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
678
-
679
- temp_dir_root : str, default: None
680
- The root directory under which `current.model.loaded` will store loaded models
681
-
682
-
683
-
561
+ times : int, default 3
562
+ Number of times to retry this task.
563
+ minutes_between_retries : int, default 2
564
+ Number of minutes between retries.
684
565
  """
685
566
  ...
686
567
 
@@ -735,6 +616,63 @@ def card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typ
735
616
  """
736
617
  ...
737
618
 
619
+ @typing.overload
620
+ 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]]]:
621
+ """
622
+ Specifies the Conda environment for the step.
623
+
624
+ Information in this decorator will augment any
625
+ attributes set in the `@conda_base` flow-level decorator. Hence,
626
+ you can use `@conda_base` to set packages required by all
627
+ steps and use `@conda` to specify step-specific overrides.
628
+
629
+ Parameters
630
+ ----------
631
+ packages : Dict[str, str], default {}
632
+ Packages to use for this step. The key is the name of the package
633
+ and the value is the version to use.
634
+ libraries : Dict[str, str], default {}
635
+ Supported for backward compatibility. When used with packages, packages will take precedence.
636
+ python : str, optional, default None
637
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
638
+ that the version used will correspond to the version of the Python interpreter used to start the run.
639
+ disabled : bool, default False
640
+ If set to True, disables @conda.
641
+ """
642
+ ...
643
+
644
+ @typing.overload
645
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
646
+ ...
647
+
648
+ @typing.overload
649
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
650
+ ...
651
+
652
+ 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):
653
+ """
654
+ Specifies the Conda environment for the step.
655
+
656
+ Information in this decorator will augment any
657
+ attributes set in the `@conda_base` flow-level decorator. Hence,
658
+ you can use `@conda_base` to set packages required by all
659
+ steps and use `@conda` to specify step-specific overrides.
660
+
661
+ Parameters
662
+ ----------
663
+ packages : Dict[str, str], default {}
664
+ Packages to use for this step. The key is the name of the package
665
+ and the value is the version to use.
666
+ libraries : Dict[str, str], default {}
667
+ Supported for backward compatibility. When used with packages, packages will take precedence.
668
+ python : str, optional, default None
669
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
670
+ that the version used will correspond to the version of the Python interpreter used to start the run.
671
+ disabled : bool, default False
672
+ If set to True, disables @conda.
673
+ """
674
+ ...
675
+
738
676
  @typing.overload
739
677
  def resources(*, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
740
678
  """
@@ -813,90 +751,137 @@ def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None]
813
751
  ...
814
752
 
815
753
  @typing.overload
816
- 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]]]:
754
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
817
755
  """
818
- Specifies the Conda environment for the step.
819
-
820
- Information in this decorator will augment any
821
- attributes set in the `@conda_base` flow-level decorator. Hence,
822
- you can use `@conda_base` to set packages required by all
823
- steps and use `@conda` to specify step-specific overrides.
824
-
825
- Parameters
756
+ Internal decorator to support Fast bakery
757
+ """
758
+ ...
759
+
760
+ @typing.overload
761
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
762
+ ...
763
+
764
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
765
+ """
766
+ Internal decorator to support Fast bakery
767
+ """
768
+ ...
769
+
770
+ @typing.overload
771
+ 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]]]:
772
+ """
773
+ Specifies a timeout for your step.
774
+
775
+ This decorator is useful if this step may hang indefinitely.
776
+
777
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
778
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
779
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
780
+
781
+ Note that all the values specified in parameters are added together so if you specify
782
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
783
+
784
+ Parameters
826
785
  ----------
827
- packages : Dict[str, str], default {}
828
- Packages to use for this step. The key is the name of the package
829
- and the value is the version to use.
830
- libraries : Dict[str, str], default {}
831
- Supported for backward compatibility. When used with packages, packages will take precedence.
832
- python : str, optional, default None
833
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
834
- that the version used will correspond to the version of the Python interpreter used to start the run.
835
- disabled : bool, default False
836
- If set to True, disables @conda.
786
+ seconds : int, default 0
787
+ Number of seconds to wait prior to timing out.
788
+ minutes : int, default 0
789
+ Number of minutes to wait prior to timing out.
790
+ hours : int, default 0
791
+ Number of hours to wait prior to timing out.
837
792
  """
838
793
  ...
839
794
 
840
795
  @typing.overload
841
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
796
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
842
797
  ...
843
798
 
844
799
  @typing.overload
845
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
800
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
846
801
  ...
847
802
 
848
- 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):
803
+ 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):
849
804
  """
850
- Specifies the Conda environment for the step.
805
+ Specifies a timeout for your step.
851
806
 
852
- Information in this decorator will augment any
853
- attributes set in the `@conda_base` flow-level decorator. Hence,
854
- you can use `@conda_base` to set packages required by all
855
- steps and use `@conda` to specify step-specific overrides.
807
+ This decorator is useful if this step may hang indefinitely.
808
+
809
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
810
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
811
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
812
+
813
+ Note that all the values specified in parameters are added together so if you specify
814
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
856
815
 
857
816
  Parameters
858
817
  ----------
859
- packages : Dict[str, str], default {}
860
- Packages to use for this step. The key is the name of the package
861
- and the value is the version to use.
862
- libraries : Dict[str, str], default {}
863
- Supported for backward compatibility. When used with packages, packages will take precedence.
864
- python : str, optional, default None
865
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
866
- that the version used will correspond to the version of the Python interpreter used to start the run.
867
- disabled : bool, default False
868
- If set to True, disables @conda.
818
+ seconds : int, default 0
819
+ Number of seconds to wait prior to timing out.
820
+ minutes : int, default 0
821
+ Number of minutes to wait prior to timing out.
822
+ hours : int, default 0
823
+ Number of hours to wait prior to timing out.
869
824
  """
870
825
  ...
871
826
 
872
827
  @typing.overload
873
- 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]]]:
828
+ 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]]]:
874
829
  """
875
- Specifies environment variables to be set prior to the execution of a step.
830
+ Enables loading / saving of models within a step.
831
+
876
832
 
877
833
  Parameters
878
834
  ----------
879
- vars : Dict[str, str], default {}
880
- Dictionary of environment variables to set.
835
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
836
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
837
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
838
+ - `current.checkpoint`
839
+ - `current.model`
840
+ - `current.huggingface_hub`
841
+
842
+ 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
843
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
844
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
845
+
846
+ temp_dir_root : str, default: None
847
+ The root directory under which `current.model.loaded` will store loaded models
848
+
849
+
850
+
881
851
  """
882
852
  ...
883
853
 
884
854
  @typing.overload
885
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
855
+ def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
886
856
  ...
887
857
 
888
858
  @typing.overload
889
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
859
+ def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
890
860
  ...
891
861
 
892
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
862
+ 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):
893
863
  """
894
- Specifies environment variables to be set prior to the execution of a step.
864
+ Enables loading / saving of models within a step.
865
+
895
866
 
896
867
  Parameters
897
868
  ----------
898
- vars : Dict[str, str], default {}
899
- Dictionary of environment variables to set.
869
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
870
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
871
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
872
+ - `current.checkpoint`
873
+ - `current.model`
874
+ - `current.huggingface_hub`
875
+
876
+ 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
877
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
878
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
879
+
880
+ temp_dir_root : str, default: None
881
+ The root directory under which `current.model.loaded` will store loaded models
882
+
883
+
884
+
900
885
  """
901
886
  ...
902
887
 
@@ -968,112 +953,6 @@ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None],
968
953
  """
969
954
  ...
970
955
 
971
- @typing.overload
972
- 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]]]:
973
- """
974
- Specifies the PyPI packages for the step.
975
-
976
- Information in this decorator will augment any
977
- attributes set in the `@pyi_base` flow-level decorator. Hence,
978
- you can use `@pypi_base` to set packages required by all
979
- steps and use `@pypi` to specify step-specific overrides.
980
-
981
- Parameters
982
- ----------
983
- packages : Dict[str, str], default: {}
984
- Packages to use for this step. The key is the name of the package
985
- and the value is the version to use.
986
- python : str, optional, default: None
987
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
988
- that the version used will correspond to the version of the Python interpreter used to start the run.
989
- """
990
- ...
991
-
992
- @typing.overload
993
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
994
- ...
995
-
996
- @typing.overload
997
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
998
- ...
999
-
1000
- 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):
1001
- """
1002
- Specifies the PyPI packages for the step.
1003
-
1004
- Information in this decorator will augment any
1005
- attributes set in the `@pyi_base` flow-level decorator. Hence,
1006
- you can use `@pypi_base` to set packages required by all
1007
- steps and use `@pypi` to specify step-specific overrides.
1008
-
1009
- Parameters
1010
- ----------
1011
- packages : Dict[str, str], default: {}
1012
- Packages to use for this step. The key is the name of the package
1013
- and the value is the version to use.
1014
- python : str, optional, default: None
1015
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1016
- that the version used will correspond to the version of the Python interpreter used to start the run.
1017
- """
1018
- ...
1019
-
1020
- @typing.overload
1021
- 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]]]:
1022
- """
1023
- Specifies a timeout for your step.
1024
-
1025
- This decorator is useful if this step may hang indefinitely.
1026
-
1027
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
1028
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
1029
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
1030
-
1031
- Note that all the values specified in parameters are added together so if you specify
1032
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
1033
-
1034
- Parameters
1035
- ----------
1036
- seconds : int, default 0
1037
- Number of seconds to wait prior to timing out.
1038
- minutes : int, default 0
1039
- Number of minutes to wait prior to timing out.
1040
- hours : int, default 0
1041
- Number of hours to wait prior to timing out.
1042
- """
1043
- ...
1044
-
1045
- @typing.overload
1046
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1047
- ...
1048
-
1049
- @typing.overload
1050
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1051
- ...
1052
-
1053
- 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):
1054
- """
1055
- Specifies a timeout for your step.
1056
-
1057
- This decorator is useful if this step may hang indefinitely.
1058
-
1059
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
1060
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
1061
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
1062
-
1063
- Note that all the values specified in parameters are added together so if you specify
1064
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
1065
-
1066
- Parameters
1067
- ----------
1068
- seconds : int, default 0
1069
- Number of seconds to wait prior to timing out.
1070
- minutes : int, default 0
1071
- Number of minutes to wait prior to timing out.
1072
- hours : int, default 0
1073
- Number of hours to wait prior to timing out.
1074
- """
1075
- ...
1076
-
1077
956
  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, hostname_resolution_timeout: int = 600) -> 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]]]:
1078
957
  """
1079
958
  Specifies that this step should execute on Kubernetes.
@@ -1143,146 +1022,232 @@ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: ty
1143
1022
  """
1144
1023
  ...
1145
1024
 
1146
- def huggingface_hub(*, temp_dir_root: 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]]]:
1025
+ @typing.overload
1026
+ 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]]]:
1147
1027
  """
1148
- Decorator that helps cache, version and store models/datasets from huggingface hub.
1028
+ Specifies the PyPI packages for the step.
1029
+
1030
+ Information in this decorator will augment any
1031
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
1032
+ you can use `@pypi_base` to set packages required by all
1033
+ steps and use `@pypi` to specify step-specific overrides.
1149
1034
 
1150
1035
  Parameters
1151
1036
  ----------
1152
- temp_dir_root : str, optional
1153
- The root directory that will hold the temporary directory where objects will be downloaded.
1154
-
1155
-
1156
- """
1157
- ...
1158
-
1159
- @typing.overload
1160
- 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]]]:
1037
+ packages : Dict[str, str], default: {}
1038
+ Packages to use for this step. The key is the name of the package
1039
+ and the value is the version to use.
1040
+ python : str, optional, default: None
1041
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1042
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1161
1043
  """
1162
- Specifies the number of times the task corresponding
1163
- to a step needs to be retried.
1044
+ ...
1045
+
1046
+ @typing.overload
1047
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1048
+ ...
1049
+
1050
+ @typing.overload
1051
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1052
+ ...
1053
+
1054
+ 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):
1055
+ """
1056
+ Specifies the PyPI packages for the step.
1164
1057
 
1165
- This decorator is useful for handling transient errors, such as networking issues.
1166
- If your task contains operations that can't be retried safely, e.g. database updates,
1167
- it is advisable to annotate it with `@retry(times=0)`.
1058
+ Information in this decorator will augment any
1059
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
1060
+ you can use `@pypi_base` to set packages required by all
1061
+ steps and use `@pypi` to specify step-specific overrides.
1168
1062
 
1169
- This can be used in conjunction with the `@catch` decorator. The `@catch`
1170
- decorator will execute a no-op task after all retries have been exhausted,
1171
- ensuring that the flow execution can continue.
1063
+ Parameters
1064
+ ----------
1065
+ packages : Dict[str, str], default: {}
1066
+ Packages to use for this step. The key is the name of the package
1067
+ and the value is the version to use.
1068
+ python : str, optional, default: None
1069
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1070
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1071
+ """
1072
+ ...
1073
+
1074
+ def huggingface_hub(*, temp_dir_root: 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]]]:
1075
+ """
1076
+ Decorator that helps cache, version and store models/datasets from huggingface hub.
1172
1077
 
1173
1078
  Parameters
1174
1079
  ----------
1175
- times : int, default 3
1176
- Number of times to retry this task.
1177
- minutes_between_retries : int, default 2
1178
- Number of minutes between retries.
1080
+ temp_dir_root : str, optional
1081
+ The root directory that will hold the temporary directory where objects will be downloaded.
1082
+
1083
+
1179
1084
  """
1180
1085
  ...
1181
1086
 
1182
1087
  @typing.overload
1183
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1088
+ 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]]]:
1089
+ """
1090
+ Enables checkpointing for a step.
1091
+
1092
+
1093
+ Parameters
1094
+ ----------
1095
+ load_policy : str, default: "fresh"
1096
+ The policy for loading the checkpoint. The following policies are supported:
1097
+ - "eager": Loads the the latest available checkpoint within the namespace.
1098
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
1099
+ will be loaded at the start of the task.
1100
+ - "none": Do not load any checkpoint
1101
+ - "fresh": Loads the lastest checkpoint created within the running Task.
1102
+ This mode helps loading checkpoints across various retry attempts of the same task.
1103
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
1104
+ created within the task will be loaded when the task is retries execution on failure.
1105
+
1106
+ temp_dir_root : str, default: None
1107
+ The root directory under which `current.checkpoint.directory` will be created.
1108
+
1109
+
1110
+
1111
+ """
1184
1112
  ...
1185
1113
 
1186
1114
  @typing.overload
1187
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1115
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1188
1116
  ...
1189
1117
 
1190
- 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):
1118
+ @typing.overload
1119
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1120
+ ...
1121
+
1122
+ 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):
1191
1123
  """
1192
- Specifies the number of times the task corresponding
1193
- to a step needs to be retried.
1194
-
1195
- This decorator is useful for handling transient errors, such as networking issues.
1196
- If your task contains operations that can't be retried safely, e.g. database updates,
1197
- it is advisable to annotate it with `@retry(times=0)`.
1124
+ Enables checkpointing for a step.
1198
1125
 
1199
- This can be used in conjunction with the `@catch` decorator. The `@catch`
1200
- decorator will execute a no-op task after all retries have been exhausted,
1201
- ensuring that the flow execution can continue.
1202
1126
 
1203
1127
  Parameters
1204
1128
  ----------
1205
- times : int, default 3
1206
- Number of times to retry this task.
1207
- minutes_between_retries : int, default 2
1208
- Number of minutes between retries.
1129
+ load_policy : str, default: "fresh"
1130
+ The policy for loading the checkpoint. The following policies are supported:
1131
+ - "eager": Loads the the latest available checkpoint within the namespace.
1132
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
1133
+ will be loaded at the start of the task.
1134
+ - "none": Do not load any checkpoint
1135
+ - "fresh": Loads the lastest checkpoint created within the running Task.
1136
+ This mode helps loading checkpoints across various retry attempts of the same task.
1137
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
1138
+ created within the task will be loaded when the task is retries execution on failure.
1139
+
1140
+ temp_dir_root : str, default: None
1141
+ The root directory under which `current.checkpoint.directory` will be created.
1142
+
1143
+
1144
+
1209
1145
  """
1210
1146
  ...
1211
1147
 
1212
- def nim(*, models: "list[NIM]", backend: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1148
+ @typing.overload
1149
+ def secrets(*, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = []) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
1213
1150
  """
1214
- This decorator is used to run NIM containers in Metaflow tasks as sidecars.
1151
+ Specifies secrets to be retrieved and injected as environment variables prior to
1152
+ the execution of a step.
1215
1153
 
1216
- User code call
1217
- -----------
1218
- @nim(
1219
- models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
1220
- backend='managed'
1221
- )
1154
+ Parameters
1155
+ ----------
1156
+ sources : List[Union[str, Dict[str, Any]]], default: []
1157
+ List of secret specs, defining how the secrets are to be retrieved
1158
+ """
1159
+ ...
1160
+
1161
+ @typing.overload
1162
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1163
+ ...
1164
+
1165
+ @typing.overload
1166
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1167
+ ...
1168
+
1169
+ def secrets(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = []):
1170
+ """
1171
+ Specifies secrets to be retrieved and injected as environment variables prior to
1172
+ the execution of a step.
1222
1173
 
1223
- Valid backend options
1224
- ---------------------
1225
- - 'managed': Outerbounds selects a compute provider based on the model.
1226
- - 🚧 'dataplane': Run in your account.
1174
+ Parameters
1175
+ ----------
1176
+ sources : List[Union[str, Dict[str, Any]]], default: []
1177
+ List of secret specs, defining how the secrets are to be retrieved
1178
+ """
1179
+ ...
1180
+
1181
+ @typing.overload
1182
+ 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]]]:
1183
+ """
1184
+ Specifies environment variables to be set prior to the execution of a step.
1227
1185
 
1228
- Valid model options
1229
- ----------------
1230
- - 'meta/llama3-8b-instruct': 8B parameter model
1231
- - 'meta/llama3-70b-instruct': 70B parameter model
1232
- - Upon request, any model here: https://nvcf.ngc.nvidia.com/functions?filter=nvidia-functions
1186
+ Parameters
1187
+ ----------
1188
+ vars : Dict[str, str], default {}
1189
+ Dictionary of environment variables to set.
1190
+ """
1191
+ ...
1192
+
1193
+ @typing.overload
1194
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1195
+ ...
1196
+
1197
+ @typing.overload
1198
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1199
+ ...
1200
+
1201
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
1202
+ """
1203
+ Specifies environment variables to be set prior to the execution of a step.
1233
1204
 
1234
1205
  Parameters
1235
1206
  ----------
1236
- models: list[NIM]
1237
- List of NIM containers running models in sidecars.
1238
- backend: str
1239
- Compute provider to run the NIM container.
1207
+ vars : Dict[str, str], default {}
1208
+ Dictionary of environment variables to set.
1240
1209
  """
1241
1210
  ...
1242
1211
 
1243
1212
  @typing.overload
1244
- 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]]:
1213
+ 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]]:
1245
1214
  """
1246
- Specifies the flow(s) that this flow depends on.
1215
+ Specifies the event(s) that this flow depends on.
1247
1216
 
1248
1217
  ```
1249
- @trigger_on_finish(flow='FooFlow')
1218
+ @trigger(event='foo')
1250
1219
  ```
1251
1220
  or
1252
1221
  ```
1253
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1222
+ @trigger(events=['foo', 'bar'])
1254
1223
  ```
1255
- This decorator respects the @project decorator and triggers the flow
1256
- when upstream runs within the same namespace complete successfully
1257
1224
 
1258
- Additionally, you can specify project aware upstream flow dependencies
1259
- by specifying the fully qualified project_flow_name.
1225
+ Additionally, you can specify the parameter mappings
1226
+ to map event payload to Metaflow parameters for the flow.
1260
1227
  ```
1261
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1228
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1262
1229
  ```
1263
1230
  or
1264
1231
  ```
1265
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1232
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1233
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1266
1234
  ```
1267
1235
 
1268
- You can also specify just the project or project branch (other values will be
1269
- inferred from the current project or project branch):
1236
+ 'parameters' can also be a list of strings and tuples like so:
1270
1237
  ```
1271
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1238
+ @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1239
+ ```
1240
+ This is equivalent to:
1241
+ ```
1242
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1272
1243
  ```
1273
-
1274
- Note that `branch` is typically one of:
1275
- - `prod`
1276
- - `user.bob`
1277
- - `test.my_experiment`
1278
- - `prod.staging`
1279
1244
 
1280
1245
  Parameters
1281
1246
  ----------
1282
- flow : Union[str, Dict[str, str]], optional, default None
1283
- Upstream flow dependency for this flow.
1284
- flows : List[Union[str, Dict[str, str]]], default []
1285
- Upstream flow dependencies for this flow.
1247
+ event : Union[str, Dict[str, Any]], optional, default None
1248
+ Event dependency for this flow.
1249
+ events : List[Union[str, Dict[str, Any]]], default []
1250
+ Events dependency for this flow.
1286
1251
  options : Dict[str, Any], default {}
1287
1252
  Backend-specific configuration for tuning eventing behavior.
1288
1253
 
@@ -1291,51 +1256,47 @@ def trigger_on_finish(*, flow: typing.Union[typing.Dict[str, str], str, None] =
1291
1256
  ...
1292
1257
 
1293
1258
  @typing.overload
1294
- def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1259
+ def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1295
1260
  ...
1296
1261
 
1297
- 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] = {}):
1262
+ 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] = {}):
1298
1263
  """
1299
- Specifies the flow(s) that this flow depends on.
1264
+ Specifies the event(s) that this flow depends on.
1300
1265
 
1301
1266
  ```
1302
- @trigger_on_finish(flow='FooFlow')
1267
+ @trigger(event='foo')
1303
1268
  ```
1304
1269
  or
1305
1270
  ```
1306
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1271
+ @trigger(events=['foo', 'bar'])
1307
1272
  ```
1308
- This decorator respects the @project decorator and triggers the flow
1309
- when upstream runs within the same namespace complete successfully
1310
1273
 
1311
- Additionally, you can specify project aware upstream flow dependencies
1312
- by specifying the fully qualified project_flow_name.
1274
+ Additionally, you can specify the parameter mappings
1275
+ to map event payload to Metaflow parameters for the flow.
1313
1276
  ```
1314
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1277
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1315
1278
  ```
1316
1279
  or
1317
1280
  ```
1318
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1281
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1282
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1319
1283
  ```
1320
1284
 
1321
- You can also specify just the project or project branch (other values will be
1322
- inferred from the current project or project branch):
1285
+ 'parameters' can also be a list of strings and tuples like so:
1323
1286
  ```
1324
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1287
+ @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1288
+ ```
1289
+ This is equivalent to:
1290
+ ```
1291
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1325
1292
  ```
1326
-
1327
- Note that `branch` is typically one of:
1328
- - `prod`
1329
- - `user.bob`
1330
- - `test.my_experiment`
1331
- - `prod.staging`
1332
1293
 
1333
1294
  Parameters
1334
1295
  ----------
1335
- flow : Union[str, Dict[str, str]], optional, default None
1336
- Upstream flow dependency for this flow.
1337
- flows : List[Union[str, Dict[str, str]]], default []
1338
- Upstream flow dependencies for this flow.
1296
+ event : Union[str, Dict[str, Any]], optional, default None
1297
+ Event dependency for this flow.
1298
+ events : List[Union[str, Dict[str, Any]]], default []
1299
+ Events dependency for this flow.
1339
1300
  options : Dict[str, Any], default {}
1340
1301
  Backend-specific configuration for tuning eventing behavior.
1341
1302
 
@@ -1343,87 +1304,6 @@ def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *
1343
1304
  """
1344
1305
  ...
1345
1306
 
1346
- 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]]:
1347
- """
1348
- 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.
1349
- This decorator only works when a flow is scheduled on Airflow and is compiled using `airflow create`. More than one `@airflow_external_task_sensor` can be added as a flow decorators. Adding more than one decorator will ensure that `start` step starts only after all sensors finish.
1350
-
1351
- Parameters
1352
- ----------
1353
- timeout : int
1354
- Time, in seconds before the task times out and fails. (Default: 3600)
1355
- poke_interval : int
1356
- Time in seconds that the job should wait in between each try. (Default: 60)
1357
- mode : str
1358
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1359
- exponential_backoff : bool
1360
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1361
- pool : str
1362
- the slot pool this task should run in,
1363
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1364
- soft_fail : bool
1365
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1366
- name : str
1367
- Name of the sensor on Airflow
1368
- description : str
1369
- Description of sensor in the Airflow UI
1370
- external_dag_id : str
1371
- The dag_id that contains the task you want to wait for.
1372
- external_task_ids : List[str]
1373
- The list of task_ids that you want to wait for.
1374
- If None (default value) the sensor waits for the DAG. (Default: None)
1375
- allowed_states : List[str]
1376
- Iterable of allowed states, (Default: ['success'])
1377
- failed_states : List[str]
1378
- Iterable of failed or dis-allowed states. (Default: None)
1379
- execution_delta : datetime.timedelta
1380
- time difference with the previous execution to look at,
1381
- the default is the same logical date as the current task or DAG. (Default: None)
1382
- check_existence: bool
1383
- Set to True to check if the external task exists or check if
1384
- the DAG to wait for exists. (Default: True)
1385
- """
1386
- ...
1387
-
1388
- @typing.overload
1389
- def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1390
- """
1391
- Specifies the PyPI packages for all steps of the flow.
1392
-
1393
- Use `@pypi_base` to set common packages required by all
1394
- steps and use `@pypi` to specify step-specific overrides.
1395
- Parameters
1396
- ----------
1397
- packages : Dict[str, str], default: {}
1398
- Packages to use for this flow. The key is the name of the package
1399
- and the value is the version to use.
1400
- python : str, optional, default: None
1401
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1402
- that the version used will correspond to the version of the Python interpreter used to start the run.
1403
- """
1404
- ...
1405
-
1406
- @typing.overload
1407
- def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1408
- ...
1409
-
1410
- def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1411
- """
1412
- Specifies the PyPI packages for all steps of the flow.
1413
-
1414
- Use `@pypi_base` to set common packages required by all
1415
- steps and use `@pypi` to specify step-specific overrides.
1416
- Parameters
1417
- ----------
1418
- packages : Dict[str, str], default: {}
1419
- Packages to use for this flow. The key is the name of the package
1420
- and the value is the version to use.
1421
- python : str, optional, default: None
1422
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1423
- that the version used will correspond to the version of the Python interpreter used to start the run.
1424
- """
1425
- ...
1426
-
1427
1307
  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]]:
1428
1308
  """
1429
1309
  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)
@@ -1466,94 +1346,109 @@ def airflow_s3_key_sensor(*, timeout: int, poke_interval: int, mode: str, expone
1466
1346
  """
1467
1347
  ...
1468
1348
 
1469
- @typing.overload
1470
- 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]]:
1349
+ 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]]:
1471
1350
  """
1472
- Specifies the Conda environment for all steps of the flow.
1473
-
1474
- Use `@conda_base` to set common libraries required by all
1475
- steps and use `@conda` to specify step-specific additions.
1351
+ 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.
1352
+ This decorator only works when a flow is scheduled on Airflow and is compiled using `airflow create`. More than one `@airflow_external_task_sensor` can be added as a flow decorators. Adding more than one decorator will ensure that `start` step starts only after all sensors finish.
1476
1353
 
1477
1354
  Parameters
1478
1355
  ----------
1479
- packages : Dict[str, str], default {}
1480
- Packages to use for this flow. The key is the name of the package
1481
- and the value is the version to use.
1482
- libraries : Dict[str, str], default {}
1483
- Supported for backward compatibility. When used with packages, packages will take precedence.
1484
- python : str, optional, default None
1485
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1486
- that the version used will correspond to the version of the Python interpreter used to start the run.
1487
- disabled : bool, default False
1488
- If set to True, disables Conda.
1356
+ timeout : int
1357
+ Time, in seconds before the task times out and fails. (Default: 3600)
1358
+ poke_interval : int
1359
+ Time in seconds that the job should wait in between each try. (Default: 60)
1360
+ mode : str
1361
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1362
+ exponential_backoff : bool
1363
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1364
+ pool : str
1365
+ the slot pool this task should run in,
1366
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1367
+ soft_fail : bool
1368
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1369
+ name : str
1370
+ Name of the sensor on Airflow
1371
+ description : str
1372
+ Description of sensor in the Airflow UI
1373
+ external_dag_id : str
1374
+ The dag_id that contains the task you want to wait for.
1375
+ external_task_ids : List[str]
1376
+ The list of task_ids that you want to wait for.
1377
+ If None (default value) the sensor waits for the DAG. (Default: None)
1378
+ allowed_states : List[str]
1379
+ Iterable of allowed states, (Default: ['success'])
1380
+ failed_states : List[str]
1381
+ Iterable of failed or dis-allowed states. (Default: None)
1382
+ execution_delta : datetime.timedelta
1383
+ time difference with the previous execution to look at,
1384
+ the default is the same logical date as the current task or DAG. (Default: None)
1385
+ check_existence: bool
1386
+ Set to True to check if the external task exists or check if
1387
+ the DAG to wait for exists. (Default: True)
1489
1388
  """
1490
1389
  ...
1491
1390
 
1492
- @typing.overload
1493
- def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1494
- ...
1495
-
1496
- def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
1391
+ def project(*, name: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1497
1392
  """
1498
- Specifies the Conda environment for all steps of the flow.
1393
+ Specifies what flows belong to the same project.
1499
1394
 
1500
- Use `@conda_base` to set common libraries required by all
1501
- steps and use `@conda` to specify step-specific additions.
1395
+ A project-specific namespace is created for all flows that
1396
+ use the same `@project(name)`.
1502
1397
 
1503
1398
  Parameters
1504
1399
  ----------
1505
- packages : Dict[str, str], default {}
1506
- Packages to use for this flow. The key is the name of the package
1507
- and the value is the version to use.
1508
- libraries : Dict[str, str], default {}
1509
- Supported for backward compatibility. When used with packages, packages will take precedence.
1510
- python : str, optional, default None
1511
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1512
- that the version used will correspond to the version of the Python interpreter used to start the run.
1513
- disabled : bool, default False
1514
- If set to True, disables Conda.
1400
+ name : str
1401
+ Project name. Make sure that the name is unique amongst all
1402
+ projects that use the same production scheduler. The name may
1403
+ contain only lowercase alphanumeric characters and underscores.
1404
+
1405
+
1515
1406
  """
1516
1407
  ...
1517
1408
 
1518
1409
  @typing.overload
1519
- 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]]:
1410
+ 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]]:
1520
1411
  """
1521
- Specifies the event(s) that this flow depends on.
1412
+ Specifies the flow(s) that this flow depends on.
1522
1413
 
1523
1414
  ```
1524
- @trigger(event='foo')
1415
+ @trigger_on_finish(flow='FooFlow')
1525
1416
  ```
1526
1417
  or
1527
1418
  ```
1528
- @trigger(events=['foo', 'bar'])
1419
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1529
1420
  ```
1421
+ This decorator respects the @project decorator and triggers the flow
1422
+ when upstream runs within the same namespace complete successfully
1530
1423
 
1531
- Additionally, you can specify the parameter mappings
1532
- to map event payload to Metaflow parameters for the flow.
1424
+ Additionally, you can specify project aware upstream flow dependencies
1425
+ by specifying the fully qualified project_flow_name.
1533
1426
  ```
1534
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1427
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1535
1428
  ```
1536
1429
  or
1537
1430
  ```
1538
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1539
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1431
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1540
1432
  ```
1541
1433
 
1542
- 'parameters' can also be a list of strings and tuples like so:
1543
- ```
1544
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1545
- ```
1546
- This is equivalent to:
1434
+ You can also specify just the project or project branch (other values will be
1435
+ inferred from the current project or project branch):
1547
1436
  ```
1548
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1437
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1549
1438
  ```
1550
1439
 
1440
+ Note that `branch` is typically one of:
1441
+ - `prod`
1442
+ - `user.bob`
1443
+ - `test.my_experiment`
1444
+ - `prod.staging`
1445
+
1551
1446
  Parameters
1552
1447
  ----------
1553
- event : Union[str, Dict[str, Any]], optional, default None
1554
- Event dependency for this flow.
1555
- events : List[Union[str, Dict[str, Any]]], default []
1556
- Events dependency for this flow.
1448
+ flow : Union[str, Dict[str, str]], optional, default None
1449
+ Upstream flow dependency for this flow.
1450
+ flows : List[Union[str, Dict[str, str]]], default []
1451
+ Upstream flow dependencies for this flow.
1557
1452
  options : Dict[str, Any], default {}
1558
1453
  Backend-specific configuration for tuning eventing behavior.
1559
1454
 
@@ -1562,47 +1457,51 @@ def trigger(*, event: typing.Union[str, typing.Dict[str, typing.Any], None] = No
1562
1457
  ...
1563
1458
 
1564
1459
  @typing.overload
1565
- def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1460
+ def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1566
1461
  ...
1567
1462
 
1568
- 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] = {}):
1463
+ 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] = {}):
1569
1464
  """
1570
- Specifies the event(s) that this flow depends on.
1465
+ Specifies the flow(s) that this flow depends on.
1571
1466
 
1572
1467
  ```
1573
- @trigger(event='foo')
1468
+ @trigger_on_finish(flow='FooFlow')
1574
1469
  ```
1575
1470
  or
1576
1471
  ```
1577
- @trigger(events=['foo', 'bar'])
1472
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1578
1473
  ```
1474
+ This decorator respects the @project decorator and triggers the flow
1475
+ when upstream runs within the same namespace complete successfully
1579
1476
 
1580
- Additionally, you can specify the parameter mappings
1581
- to map event payload to Metaflow parameters for the flow.
1477
+ Additionally, you can specify project aware upstream flow dependencies
1478
+ by specifying the fully qualified project_flow_name.
1582
1479
  ```
1583
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1480
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1584
1481
  ```
1585
1482
  or
1586
1483
  ```
1587
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1588
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1484
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1589
1485
  ```
1590
1486
 
1591
- 'parameters' can also be a list of strings and tuples like so:
1592
- ```
1593
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1594
- ```
1595
- This is equivalent to:
1487
+ You can also specify just the project or project branch (other values will be
1488
+ inferred from the current project or project branch):
1596
1489
  ```
1597
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1490
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1598
1491
  ```
1599
1492
 
1493
+ Note that `branch` is typically one of:
1494
+ - `prod`
1495
+ - `user.bob`
1496
+ - `test.my_experiment`
1497
+ - `prod.staging`
1498
+
1600
1499
  Parameters
1601
1500
  ----------
1602
- event : Union[str, Dict[str, Any]], optional, default None
1603
- Event dependency for this flow.
1604
- events : List[Union[str, Dict[str, Any]]], default []
1605
- Events dependency for this flow.
1501
+ flow : Union[str, Dict[str, str]], optional, default None
1502
+ Upstream flow dependency for this flow.
1503
+ flows : List[Union[str, Dict[str, str]]], default []
1504
+ Upstream flow dependencies for this flow.
1606
1505
  options : Dict[str, Any], default {}
1607
1506
  Backend-specific configuration for tuning eventing behavior.
1608
1507
 
@@ -1610,24 +1509,6 @@ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: t
1610
1509
  """
1611
1510
  ...
1612
1511
 
1613
- def project(*, name: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1614
- """
1615
- Specifies what flows belong to the same project.
1616
-
1617
- A project-specific namespace is created for all flows that
1618
- use the same `@project(name)`.
1619
-
1620
- Parameters
1621
- ----------
1622
- name : str
1623
- Project name. Make sure that the name is unique amongst all
1624
- projects that use the same production scheduler. The name may
1625
- contain only lowercase alphanumeric characters and underscores.
1626
-
1627
-
1628
- """
1629
- ...
1630
-
1631
1512
  @typing.overload
1632
1513
  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]]:
1633
1514
  """
@@ -1677,6 +1558,125 @@ def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly:
1677
1558
  """
1678
1559
  ...
1679
1560
 
1561
+ @typing.overload
1562
+ 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]]:
1563
+ """
1564
+ Specifies the Conda environment for all steps of the flow.
1565
+
1566
+ Use `@conda_base` to set common libraries required by all
1567
+ steps and use `@conda` to specify step-specific additions.
1568
+
1569
+ Parameters
1570
+ ----------
1571
+ packages : Dict[str, str], default {}
1572
+ Packages to use for this flow. The key is the name of the package
1573
+ and the value is the version to use.
1574
+ libraries : Dict[str, str], default {}
1575
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1576
+ python : str, optional, default None
1577
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1578
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1579
+ disabled : bool, default False
1580
+ If set to True, disables Conda.
1581
+ """
1582
+ ...
1583
+
1584
+ @typing.overload
1585
+ def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1586
+ ...
1587
+
1588
+ def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
1589
+ """
1590
+ Specifies the Conda environment for all steps of the flow.
1591
+
1592
+ Use `@conda_base` to set common libraries required by all
1593
+ steps and use `@conda` to specify step-specific additions.
1594
+
1595
+ Parameters
1596
+ ----------
1597
+ packages : Dict[str, str], default {}
1598
+ Packages to use for this flow. The key is the name of the package
1599
+ and the value is the version to use.
1600
+ libraries : Dict[str, str], default {}
1601
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1602
+ python : str, optional, default None
1603
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1604
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1605
+ disabled : bool, default False
1606
+ If set to True, disables Conda.
1607
+ """
1608
+ ...
1609
+
1610
+ @typing.overload
1611
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1612
+ """
1613
+ Specifies the PyPI packages for all steps of the flow.
1614
+
1615
+ Use `@pypi_base` to set common packages required by all
1616
+ steps and use `@pypi` to specify step-specific overrides.
1617
+ Parameters
1618
+ ----------
1619
+ packages : Dict[str, str], default: {}
1620
+ Packages to use for this flow. The key is the name of the package
1621
+ and the value is the version to use.
1622
+ python : str, optional, default: None
1623
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1624
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1625
+ """
1626
+ ...
1627
+
1628
+ @typing.overload
1629
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1630
+ ...
1631
+
1632
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1633
+ """
1634
+ Specifies the PyPI packages for all steps of the flow.
1635
+
1636
+ Use `@pypi_base` to set common packages required by all
1637
+ steps and use `@pypi` to specify step-specific overrides.
1638
+ Parameters
1639
+ ----------
1640
+ packages : Dict[str, str], default: {}
1641
+ Packages to use for this flow. The key is the name of the package
1642
+ and the value is the version to use.
1643
+ python : str, optional, default: None
1644
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1645
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1646
+ """
1647
+ ...
1648
+
1649
+ def nim(*, models: "list[NIM]", backend: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1650
+ """
1651
+ This decorator is used to run NIM containers in Metaflow tasks as sidecars.
1652
+
1653
+ User code call
1654
+ -----------
1655
+ @nim(
1656
+ models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
1657
+ backend='managed'
1658
+ )
1659
+
1660
+ Valid backend options
1661
+ ---------------------
1662
+ - 'managed': Outerbounds selects a compute provider based on the model.
1663
+ - 🚧 'dataplane': Run in your account.
1664
+
1665
+ Valid model options
1666
+ ----------------
1667
+ - 'meta/llama3-8b-instruct': 8B parameter model
1668
+ - 'meta/llama3-70b-instruct': 70B parameter model
1669
+ - Upon request, any model here: https://nvcf.ngc.nvidia.com/functions?filter=nvidia-functions
1670
+
1671
+ Parameters
1672
+ ----------
1673
+ models: list[NIM]
1674
+ List of NIM containers running models in sidecars.
1675
+ backend: str
1676
+ Compute provider to run the NIM container.
1677
+ """
1678
+ ...
1679
+
1680
1680
  def namespace(ns: typing.Optional[str]) -> typing.Optional[str]:
1681
1681
  """
1682
1682
  Switch namespace to the one provided.