ob-metaflow-stubs 6.0.3.112__py2.py3-none-any.whl → 6.0.3.113__py2.py3-none-any.whl

Sign up to get free protection for your applications and to get access to all the features.
Files changed (140) hide show
  1. metaflow-stubs/__init__.pyi +501 -501
  2. metaflow-stubs/cards.pyi +5 -5
  3. metaflow-stubs/cli.pyi +3 -3
  4. metaflow-stubs/client/__init__.pyi +4 -4
  5. metaflow-stubs/client/core.pyi +7 -7
  6. metaflow-stubs/client/filecache.pyi +3 -3
  7. metaflow-stubs/clone_util.pyi +2 -2
  8. metaflow-stubs/events.pyi +3 -3
  9. metaflow-stubs/exception.pyi +2 -2
  10. metaflow-stubs/flowspec.pyi +5 -5
  11. metaflow-stubs/generated_for.txt +1 -1
  12. metaflow-stubs/includefile.pyi +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 +92 -92
  18. metaflow-stubs/mflog/mflog.pyi +2 -2
  19. metaflow-stubs/multicore_utils.pyi +2 -2
  20. metaflow-stubs/parameters.pyi +4 -4
  21. metaflow-stubs/plugins/__init__.pyi +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 +3 -3
  27. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +4 -4
  28. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +4 -4
  29. metaflow-stubs/plugins/argo/__init__.pyi +2 -2
  30. metaflow-stubs/plugins/argo/argo_client.pyi +2 -2
  31. metaflow-stubs/plugins/argo/argo_events.pyi +2 -2
  32. metaflow-stubs/plugins/argo/argo_workflows.pyi +5 -5
  33. metaflow-stubs/plugins/argo/argo_workflows_cli.pyi +7 -7
  34. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +6 -6
  35. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +3 -3
  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 +4 -4
  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 +4 -4
  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 +5 -5
  61. metaflow-stubs/plugins/cards/card_client.pyi +3 -3
  62. metaflow-stubs/plugins/cards/card_creator.pyi +2 -2
  63. metaflow-stubs/plugins/cards/card_datastore.pyi +2 -2
  64. metaflow-stubs/plugins/cards/card_decorator.pyi +2 -2
  65. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +2 -2
  66. metaflow-stubs/plugins/cards/card_modules/basic.pyi +3 -3
  67. metaflow-stubs/plugins/cards/card_modules/card.pyi +2 -2
  68. metaflow-stubs/plugins/cards/card_modules/components.pyi +4 -4
  69. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +2 -2
  70. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +2 -2
  71. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +2 -2
  72. metaflow-stubs/plugins/cards/card_resolver.pyi +2 -2
  73. metaflow-stubs/plugins/cards/component_serializer.pyi +4 -4
  74. metaflow-stubs/plugins/cards/exception.pyi +2 -2
  75. metaflow-stubs/plugins/catch_decorator.pyi +3 -3
  76. metaflow-stubs/plugins/datatools/__init__.pyi +4 -4
  77. metaflow-stubs/plugins/datatools/local.pyi +2 -2
  78. metaflow-stubs/plugins/datatools/s3/__init__.pyi +4 -4
  79. metaflow-stubs/plugins/datatools/s3/s3.pyi +4 -4
  80. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +2 -2
  81. metaflow-stubs/plugins/datatools/s3/s3util.pyi +2 -2
  82. metaflow-stubs/plugins/debug_logger.pyi +2 -2
  83. metaflow-stubs/plugins/debug_monitor.pyi +2 -2
  84. metaflow-stubs/plugins/environment_decorator.pyi +2 -2
  85. metaflow-stubs/plugins/events_decorator.pyi +2 -2
  86. metaflow-stubs/plugins/frameworks/__init__.pyi +2 -2
  87. metaflow-stubs/plugins/frameworks/pytorch.pyi +2 -2
  88. metaflow-stubs/plugins/gcp/__init__.pyi +2 -2
  89. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +4 -4
  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 +4 -4
  98. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +2 -2
  99. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +3 -3
  100. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +2 -2
  101. metaflow-stubs/plugins/logs_cli.pyi +4 -4
  102. metaflow-stubs/plugins/package_cli.pyi +2 -2
  103. metaflow-stubs/plugins/parallel_decorator.pyi +2 -2
  104. metaflow-stubs/plugins/perimeters.pyi +2 -2
  105. metaflow-stubs/plugins/project_decorator.pyi +2 -2
  106. metaflow-stubs/plugins/pypi/__init__.pyi +2 -2
  107. metaflow-stubs/plugins/pypi/conda_decorator.pyi +2 -2
  108. metaflow-stubs/plugins/pypi/conda_environment.pyi +4 -4
  109. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +2 -2
  110. metaflow-stubs/plugins/pypi/pypi_environment.pyi +2 -2
  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 +5 -5
  119. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +2 -2
  120. metaflow-stubs/plugins/timeout_decorator.pyi +3 -3
  121. metaflow-stubs/procpoll.pyi +2 -2
  122. metaflow-stubs/profilers/__init__.pyi +2 -2
  123. metaflow-stubs/pylint_wrapper.pyi +2 -2
  124. metaflow-stubs/runner/__init__.pyi +2 -2
  125. metaflow-stubs/runner/deployer.pyi +4 -4
  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 +3 -3
  133. metaflow-stubs/system/system_monitor.pyi +3 -3
  134. metaflow-stubs/tagging_util.pyi +2 -2
  135. metaflow-stubs/tuple_util.pyi +2 -2
  136. {ob_metaflow_stubs-6.0.3.112.dist-info → ob_metaflow_stubs-6.0.3.113.dist-info}/METADATA +1 -1
  137. ob_metaflow_stubs-6.0.3.113.dist-info/RECORD +140 -0
  138. ob_metaflow_stubs-6.0.3.112.dist-info/RECORD +0 -140
  139. {ob_metaflow_stubs-6.0.3.112.dist-info → ob_metaflow_stubs-6.0.3.113.dist-info}/WHEEL +0 -0
  140. {ob_metaflow_stubs-6.0.3.112.dist-info → ob_metaflow_stubs-6.0.3.113.dist-info}/top_level.txt +0 -0
@@ -1,24 +1,24 @@
1
1
  ##################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
- # MF version: 2.12.25.2+obcheckpoint(0.1.0);ob(v1) #
4
- # Generated on 2024-10-21T21:22:37.544481 #
3
+ # MF version: 2.12.25.2+obcheckpoint(0.1.1);ob(v1) #
4
+ # Generated on 2024-10-22T22:37:22.374436 #
5
5
  ##################################################################################
6
6
 
7
7
  from __future__ import annotations
8
8
 
9
9
  import typing
10
10
  if typing.TYPE_CHECKING:
11
- import typing
11
+ import datetime
12
+ import metaflow_extensions.obcheckpoint.plugins.machine_learning_utilities.datastructures
13
+ import metaflow.datastore.inputs
14
+ import metaflow.parameters
12
15
  import metaflow.flowspec
16
+ import typing
17
+ import metaflow._vendor.click.types
13
18
  import metaflow.client.core
14
- import metaflow.events
15
- import metaflow_extensions.obcheckpoint.plugins.machine_learning_utilities.datastructures
16
19
  import metaflow.metaflow_current
17
20
  import metaflow.runner.metaflow_runner
18
- import metaflow.datastore.inputs
19
- import metaflow._vendor.click.types
20
- import metaflow.parameters
21
- import datetime
21
+ import metaflow.events
22
22
  FlowSpecDerived = typing.TypeVar("FlowSpecDerived", bound="FlowSpec", contravariant=False, covariant=False)
23
23
  StepFlag = typing.NewType("StepFlag", bool)
24
24
 
@@ -512,159 +512,103 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
512
512
  """
513
513
  ...
514
514
 
515
- def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: typing.Optional[str] = None, image_pull_policy: str = "KUBERNETES_IMAGE_PULL_POLICY", service_account: str = "METAFLOW_KUBERNETES_SERVICE_ACCOUNT", secrets: typing.Optional[typing.List[str]] = None, node_selector: typing.Union[typing.Dict[str, str], str, None] = None, namespace: str = "METAFLOW_KUBERNETES_NAMESPACE", gpu: typing.Optional[int] = None, gpu_vendor: str = "KUBERNETES_GPU_VENDOR", tolerations: typing.List[str] = [], use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = "/metaflow_temp", persistent_volume_claims: typing.Optional[typing.Dict[str, str]] = None, shared_memory: typing.Optional[int] = None, port: typing.Optional[int] = None, compute_pool: typing.Optional[str] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
515
+ @typing.overload
516
+ 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]]]:
516
517
  """
517
- Specifies that this step should execute on Kubernetes.
518
+ Specifies that the step will success under all circumstances.
519
+
520
+ The decorator will create an optional artifact, specified by `var`, which
521
+ contains the exception raised. You can use it to detect the presence
522
+ of errors, indicating that all happy-path artifacts produced by the step
523
+ are missing.
518
524
 
519
525
  Parameters
520
526
  ----------
521
- cpu : int, default 1
522
- Number of CPUs required for this step. If `@resources` is
523
- also present, the maximum value from all decorators is used.
524
- memory : int, default 4096
525
- Memory size (in MB) required for this step. If
526
- `@resources` is also present, the maximum value from all decorators is
527
- used.
528
- disk : int, default 10240
529
- Disk size (in MB) required for this step. If
530
- `@resources` is also present, the maximum value from all decorators is
531
- used.
532
- image : str, optional, default None
533
- Docker image to use when launching on Kubernetes. If not specified, and
534
- METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
535
- not, a default Docker image mapping to the current version of Python is used.
536
- image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
537
- If given, the imagePullPolicy to be applied to the Docker image of the step.
538
- service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
539
- Kubernetes service account to use when launching pod in Kubernetes.
540
- secrets : List[str], optional, default None
541
- Kubernetes secrets to use when launching pod in Kubernetes. These
542
- secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
543
- in Metaflow configuration.
544
- node_selector: Union[Dict[str,str], str], optional, default None
545
- Kubernetes node selector(s) to apply to the pod running the task.
546
- Can be passed in as a comma separated string of values e.g. "kubernetes.io/os=linux,kubernetes.io/arch=amd64"
547
- or as a dictionary {"kubernetes.io/os": "linux", "kubernetes.io/arch": "amd64"}
548
- namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
549
- Kubernetes namespace to use when launching pod in Kubernetes.
550
- gpu : int, optional, default None
551
- Number of GPUs required for this step. A value of zero implies that
552
- the scheduled node should not have GPUs.
553
- gpu_vendor : str, default KUBERNETES_GPU_VENDOR
554
- The vendor of the GPUs to be used for this step.
555
- tolerations : List[str], default []
556
- The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
557
- Kubernetes tolerations to use when launching pod in Kubernetes.
558
- use_tmpfs : bool, default False
559
- This enables an explicit tmpfs mount for this step.
560
- tmpfs_tempdir : bool, default True
561
- sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
562
- tmpfs_size : int, optional, default: None
563
- The value for the size (in MiB) of the tmpfs mount for this step.
564
- This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
565
- memory allocated for this step.
566
- tmpfs_path : str, optional, default /metaflow_temp
567
- Path to tmpfs mount for this step.
568
- persistent_volume_claims : Dict[str, str], optional, default None
569
- A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
570
- volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
571
- shared_memory: int, optional
572
- Shared memory size (in MiB) required for this step
573
- port: int, optional
574
- Port number to specify in the Kubernetes job object
575
- compute_pool : str, optional, default None
576
- Compute pool to be used for for this step.
577
- If not specified, any accessible compute pool within the perimeter is used.
527
+ var : str, optional, default None
528
+ Name of the artifact in which to store the caught exception.
529
+ If not specified, the exception is not stored.
530
+ print_exception : bool, default True
531
+ Determines whether or not the exception is printed to
532
+ stdout when caught.
578
533
  """
579
534
  ...
580
535
 
581
- 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]]]:
536
+ @typing.overload
537
+ def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
538
+ ...
539
+
540
+ @typing.overload
541
+ def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
542
+ ...
543
+
544
+ def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, var: typing.Optional[str] = None, print_exception: bool = True):
582
545
  """
583
- Decorator that helps cache, version and store models/datasets from huggingface hub.
546
+ Specifies that the step will success under all circumstances.
547
+
548
+ The decorator will create an optional artifact, specified by `var`, which
549
+ contains the exception raised. You can use it to detect the presence
550
+ of errors, indicating that all happy-path artifacts produced by the step
551
+ are missing.
584
552
 
585
553
  Parameters
586
554
  ----------
587
- temp_dir_root : str, optional
588
- The root directory that will hold the temporary directory where objects will be downloaded.
589
-
590
-
555
+ var : str, optional, default None
556
+ Name of the artifact in which to store the caught exception.
557
+ If not specified, the exception is not stored.
558
+ print_exception : bool, default True
559
+ Determines whether or not the exception is printed to
560
+ stdout when caught.
591
561
  """
592
562
  ...
593
563
 
594
564
  @typing.overload
595
- def resources(*, cpu: int = 1, gpu: int = 0, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
565
+ 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]]]:
596
566
  """
597
- Specifies the resources needed when executing this step.
598
-
599
- Use `@resources` to specify the resource requirements
600
- independently of the specific compute layer (`@batch`, `@kubernetes`).
567
+ Creates a human-readable report, a Metaflow Card, after this step completes.
601
568
 
602
- You can choose the compute layer on the command line by executing e.g.
603
- ```
604
- python myflow.py run --with batch
605
- ```
606
- or
607
- ```
608
- python myflow.py run --with kubernetes
609
- ```
610
- which executes the flow on the desired system using the
611
- requirements specified in `@resources`.
569
+ Note that you may add multiple `@card` decorators in a step with different parameters.
612
570
 
613
571
  Parameters
614
572
  ----------
615
- cpu : int, default 1
616
- Number of CPUs required for this step.
617
- gpu : int, default 0
618
- Number of GPUs required for this step.
619
- disk : int, optional, default None
620
- Disk size (in MB) required for this step. Only applies on Kubernetes.
621
- memory : int, default 4096
622
- Memory size (in MB) required for this step.
623
- shared_memory : int, optional, default None
624
- The value for the size (in MiB) of the /dev/shm volume for this step.
625
- This parameter maps to the `--shm-size` option in Docker.
573
+ type : str, default 'default'
574
+ Card type.
575
+ id : str, optional, default None
576
+ If multiple cards are present, use this id to identify this card.
577
+ options : Dict[str, Any], default {}
578
+ Options passed to the card. The contents depend on the card type.
579
+ timeout : int, default 45
580
+ Interrupt reporting if it takes more than this many seconds.
581
+
582
+
626
583
  """
627
584
  ...
628
585
 
629
586
  @typing.overload
630
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
587
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
631
588
  ...
632
589
 
633
590
  @typing.overload
634
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
591
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
635
592
  ...
636
593
 
637
- def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, cpu: int = 1, gpu: int = 0, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None):
594
+ 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):
638
595
  """
639
- Specifies the resources needed when executing this step.
640
-
641
- Use `@resources` to specify the resource requirements
642
- independently of the specific compute layer (`@batch`, `@kubernetes`).
596
+ Creates a human-readable report, a Metaflow Card, after this step completes.
643
597
 
644
- You can choose the compute layer on the command line by executing e.g.
645
- ```
646
- python myflow.py run --with batch
647
- ```
648
- or
649
- ```
650
- python myflow.py run --with kubernetes
651
- ```
652
- which executes the flow on the desired system using the
653
- requirements specified in `@resources`.
598
+ Note that you may add multiple `@card` decorators in a step with different parameters.
654
599
 
655
600
  Parameters
656
601
  ----------
657
- cpu : int, default 1
658
- Number of CPUs required for this step.
659
- gpu : int, default 0
660
- Number of GPUs required for this step.
661
- disk : int, optional, default None
662
- Disk size (in MB) required for this step. Only applies on Kubernetes.
663
- memory : int, default 4096
664
- Memory size (in MB) required for this step.
665
- shared_memory : int, optional, default None
666
- The value for the size (in MiB) of the /dev/shm volume for this step.
667
- This parameter maps to the `--shm-size` option in Docker.
602
+ type : str, default 'default'
603
+ Card type.
604
+ id : str, optional, default None
605
+ If multiple cards are present, use this id to identify this card.
606
+ options : Dict[str, Any], default {}
607
+ Options passed to the card. The contents depend on the card type.
608
+ timeout : int, default 45
609
+ Interrupt reporting if it takes more than this many seconds.
610
+
611
+
668
612
  """
669
613
  ...
670
614
 
@@ -722,51 +666,80 @@ def retry(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
722
666
  ...
723
667
 
724
668
  @typing.overload
725
- 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]]]:
669
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
726
670
  """
727
- Creates a human-readable report, a Metaflow Card, after this step completes.
671
+ Decorator prototype for all step decorators. This function gets specialized
672
+ and imported for all decorators types by _import_plugin_decorators().
673
+ """
674
+ ...
675
+
676
+ @typing.overload
677
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
678
+ ...
679
+
680
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
681
+ """
682
+ Decorator prototype for all step decorators. This function gets specialized
683
+ and imported for all decorators types by _import_plugin_decorators().
684
+ """
685
+ ...
686
+
687
+ @typing.overload
688
+ 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]]]:
689
+ """
690
+ Enables loading / saving of models within a step.
728
691
 
729
- Note that you may add multiple `@card` decorators in a step with different parameters.
730
692
 
731
693
  Parameters
732
694
  ----------
733
- type : str, default 'default'
734
- Card type.
735
- id : str, optional, default None
736
- If multiple cards are present, use this id to identify this card.
737
- options : Dict[str, Any], default {}
738
- Options passed to the card. The contents depend on the card type.
739
- timeout : int, default 45
740
- Interrupt reporting if it takes more than this many seconds.
695
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
696
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
697
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
698
+ - `current.checkpoint`
699
+ - `current.model`
700
+ - `current.huggingface_hub`
701
+
702
+ 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
703
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
704
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
705
+
706
+ temp_dir_root : str, default: None
707
+ The root directory under which `current.model.loaded` will store loaded models
708
+
741
709
 
742
710
 
743
711
  """
744
712
  ...
745
713
 
746
714
  @typing.overload
747
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
715
+ def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
748
716
  ...
749
717
 
750
718
  @typing.overload
751
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
719
+ def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
752
720
  ...
753
721
 
754
- 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):
722
+ 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):
755
723
  """
756
- Creates a human-readable report, a Metaflow Card, after this step completes.
724
+ Enables loading / saving of models within a step.
757
725
 
758
- Note that you may add multiple `@card` decorators in a step with different parameters.
759
726
 
760
727
  Parameters
761
728
  ----------
762
- type : str, default 'default'
763
- Card type.
764
- id : str, optional, default None
765
- If multiple cards are present, use this id to identify this card.
766
- options : Dict[str, Any], default {}
767
- Options passed to the card. The contents depend on the card type.
768
- timeout : int, default 45
769
- Interrupt reporting if it takes more than this many seconds.
729
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
730
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
731
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
732
+ - `current.checkpoint`
733
+ - `current.model`
734
+ - `current.huggingface_hub`
735
+
736
+ 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
737
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
738
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
739
+
740
+ temp_dir_root : str, default: None
741
+ The root directory under which `current.model.loaded` will store loaded models
742
+
770
743
 
771
744
 
772
745
  """
@@ -822,21 +795,35 @@ def pypi(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typ
822
795
  ...
823
796
 
824
797
  @typing.overload
825
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
798
+ 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]]]:
826
799
  """
827
- Decorator prototype for all step decorators. This function gets specialized
828
- and imported for all decorators types by _import_plugin_decorators().
800
+ Specifies secrets to be retrieved and injected as environment variables prior to
801
+ the execution of a step.
802
+
803
+ Parameters
804
+ ----------
805
+ sources : List[Union[str, Dict[str, Any]]], default: []
806
+ List of secret specs, defining how the secrets are to be retrieved
829
807
  """
830
808
  ...
831
809
 
832
810
  @typing.overload
833
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
811
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
834
812
  ...
835
813
 
836
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
814
+ @typing.overload
815
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
816
+ ...
817
+
818
+ 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]]] = []):
837
819
  """
838
- Decorator prototype for all step decorators. This function gets specialized
839
- and imported for all decorators types by _import_plugin_decorators().
820
+ Specifies secrets to be retrieved and injected as environment variables prior to
821
+ the execution of a step.
822
+
823
+ Parameters
824
+ ----------
825
+ sources : List[Union[str, Dict[str, Any]]], default: []
826
+ List of secret specs, defining how the secrets are to be retrieved
840
827
  """
841
828
  ...
842
829
 
@@ -897,116 +884,157 @@ def timeout(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None],
897
884
  """
898
885
  ...
899
886
 
900
- @typing.overload
901
- 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]]]:
902
- """
903
- Specifies that the step will success under all circumstances.
904
-
905
- The decorator will create an optional artifact, specified by `var`, which
906
- contains the exception raised. You can use it to detect the presence
907
- of errors, indicating that all happy-path artifacts produced by the step
908
- are missing.
909
-
910
- Parameters
911
- ----------
912
- var : str, optional, default None
913
- Name of the artifact in which to store the caught exception.
914
- If not specified, the exception is not stored.
915
- print_exception : bool, default True
916
- Determines whether or not the exception is printed to
917
- stdout when caught.
918
- """
919
- ...
920
-
921
- @typing.overload
922
- def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
923
- ...
924
-
925
- @typing.overload
926
- def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
927
- ...
928
-
929
- def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, var: typing.Optional[str] = None, print_exception: bool = True):
887
+ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: typing.Optional[str] = None, image_pull_policy: str = "KUBERNETES_IMAGE_PULL_POLICY", service_account: str = "METAFLOW_KUBERNETES_SERVICE_ACCOUNT", secrets: typing.Optional[typing.List[str]] = None, node_selector: typing.Union[typing.Dict[str, str], str, None] = None, namespace: str = "METAFLOW_KUBERNETES_NAMESPACE", gpu: typing.Optional[int] = None, gpu_vendor: str = "KUBERNETES_GPU_VENDOR", tolerations: typing.List[str] = [], use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = "/metaflow_temp", persistent_volume_claims: typing.Optional[typing.Dict[str, str]] = None, shared_memory: typing.Optional[int] = None, port: typing.Optional[int] = None, compute_pool: typing.Optional[str] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
930
888
  """
931
- Specifies that the step will success under all circumstances.
932
-
933
- The decorator will create an optional artifact, specified by `var`, which
934
- contains the exception raised. You can use it to detect the presence
935
- of errors, indicating that all happy-path artifacts produced by the step
936
- are missing.
889
+ Specifies that this step should execute on Kubernetes.
937
890
 
938
891
  Parameters
939
892
  ----------
940
- var : str, optional, default None
941
- Name of the artifact in which to store the caught exception.
942
- If not specified, the exception is not stored.
943
- print_exception : bool, default True
944
- Determines whether or not the exception is printed to
945
- stdout when caught.
893
+ cpu : int, default 1
894
+ Number of CPUs required for this step. If `@resources` is
895
+ also present, the maximum value from all decorators is used.
896
+ memory : int, default 4096
897
+ Memory size (in MB) required for this step. If
898
+ `@resources` is also present, the maximum value from all decorators is
899
+ used.
900
+ disk : int, default 10240
901
+ Disk size (in MB) required for this step. If
902
+ `@resources` is also present, the maximum value from all decorators is
903
+ used.
904
+ image : str, optional, default None
905
+ Docker image to use when launching on Kubernetes. If not specified, and
906
+ METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
907
+ not, a default Docker image mapping to the current version of Python is used.
908
+ image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
909
+ If given, the imagePullPolicy to be applied to the Docker image of the step.
910
+ service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
911
+ Kubernetes service account to use when launching pod in Kubernetes.
912
+ secrets : List[str], optional, default None
913
+ Kubernetes secrets to use when launching pod in Kubernetes. These
914
+ secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
915
+ in Metaflow configuration.
916
+ node_selector: Union[Dict[str,str], str], optional, default None
917
+ Kubernetes node selector(s) to apply to the pod running the task.
918
+ Can be passed in as a comma separated string of values e.g. "kubernetes.io/os=linux,kubernetes.io/arch=amd64"
919
+ or as a dictionary {"kubernetes.io/os": "linux", "kubernetes.io/arch": "amd64"}
920
+ namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
921
+ Kubernetes namespace to use when launching pod in Kubernetes.
922
+ gpu : int, optional, default None
923
+ Number of GPUs required for this step. A value of zero implies that
924
+ the scheduled node should not have GPUs.
925
+ gpu_vendor : str, default KUBERNETES_GPU_VENDOR
926
+ The vendor of the GPUs to be used for this step.
927
+ tolerations : List[str], default []
928
+ The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
929
+ Kubernetes tolerations to use when launching pod in Kubernetes.
930
+ use_tmpfs : bool, default False
931
+ This enables an explicit tmpfs mount for this step.
932
+ tmpfs_tempdir : bool, default True
933
+ sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
934
+ tmpfs_size : int, optional, default: None
935
+ The value for the size (in MiB) of the tmpfs mount for this step.
936
+ This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
937
+ memory allocated for this step.
938
+ tmpfs_path : str, optional, default /metaflow_temp
939
+ Path to tmpfs mount for this step.
940
+ persistent_volume_claims : Dict[str, str], optional, default None
941
+ A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
942
+ volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
943
+ shared_memory: int, optional
944
+ Shared memory size (in MiB) required for this step
945
+ port: int, optional
946
+ Port number to specify in the Kubernetes job object
947
+ compute_pool : str, optional, default None
948
+ Compute pool to be used for for this step.
949
+ If not specified, any accessible compute pool within the perimeter is used.
946
950
  """
947
951
  ...
948
952
 
949
953
  @typing.overload
950
- 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]]]:
954
+ 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]]]:
951
955
  """
952
- Specifies secrets to be retrieved and injected as environment variables prior to
953
- the execution of a step.
956
+ Specifies environment variables to be set prior to the execution of a step.
954
957
 
955
958
  Parameters
956
959
  ----------
957
- sources : List[Union[str, Dict[str, Any]]], default: []
958
- List of secret specs, defining how the secrets are to be retrieved
960
+ vars : Dict[str, str], default {}
961
+ Dictionary of environment variables to set.
959
962
  """
960
963
  ...
961
964
 
962
965
  @typing.overload
963
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
966
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
964
967
  ...
965
968
 
966
969
  @typing.overload
967
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
970
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
968
971
  ...
969
972
 
970
- 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]]] = []):
973
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
971
974
  """
972
- Specifies secrets to be retrieved and injected as environment variables prior to
973
- the execution of a step.
975
+ Specifies environment variables to be set prior to the execution of a step.
974
976
 
975
977
  Parameters
976
978
  ----------
977
- sources : List[Union[str, Dict[str, Any]]], default: []
978
- List of secret specs, defining how the secrets are to be retrieved
979
+ vars : Dict[str, str], default {}
980
+ Dictionary of environment variables to set.
979
981
  """
980
982
  ...
981
983
 
982
984
  @typing.overload
983
- 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]]]:
985
+ 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]]]:
984
986
  """
985
- Specifies environment variables to be set prior to the execution of a step.
987
+ Specifies the Conda environment for the step.
988
+
989
+ Information in this decorator will augment any
990
+ attributes set in the `@conda_base` flow-level decorator. Hence,
991
+ you can use `@conda_base` to set packages required by all
992
+ steps and use `@conda` to specify step-specific overrides.
986
993
 
987
994
  Parameters
988
995
  ----------
989
- vars : Dict[str, str], default {}
990
- Dictionary of environment variables to set.
996
+ packages : Dict[str, str], default {}
997
+ Packages to use for this step. The key is the name of the package
998
+ and the value is the version to use.
999
+ libraries : Dict[str, str], default {}
1000
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1001
+ python : str, optional, default None
1002
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1003
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1004
+ disabled : bool, default False
1005
+ If set to True, disables @conda.
991
1006
  """
992
1007
  ...
993
1008
 
994
1009
  @typing.overload
995
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1010
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
996
1011
  ...
997
1012
 
998
1013
  @typing.overload
999
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1014
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1000
1015
  ...
1001
1016
 
1002
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
1017
+ 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):
1003
1018
  """
1004
- Specifies environment variables to be set prior to the execution of a step.
1019
+ Specifies the Conda environment for the step.
1020
+
1021
+ Information in this decorator will augment any
1022
+ attributes set in the `@conda_base` flow-level decorator. Hence,
1023
+ you can use `@conda_base` to set packages required by all
1024
+ steps and use `@conda` to specify step-specific overrides.
1005
1025
 
1006
1026
  Parameters
1007
1027
  ----------
1008
- vars : Dict[str, str], default {}
1009
- Dictionary of environment variables to set.
1028
+ packages : Dict[str, str], default {}
1029
+ Packages to use for this step. The key is the name of the package
1030
+ and the value is the version to use.
1031
+ libraries : Dict[str, str], default {}
1032
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1033
+ python : str, optional, default None
1034
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1035
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1036
+ disabled : bool, default False
1037
+ If set to True, disables @conda.
1010
1038
  """
1011
1039
  ...
1012
1040
 
@@ -1072,137 +1100,109 @@ def checkpoint(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None
1072
1100
  ...
1073
1101
 
1074
1102
  @typing.overload
1075
- 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]]]:
1103
+ def resources(*, cpu: int = 1, gpu: int = 0, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
1076
1104
  """
1077
- Enables loading / saving of models within a step.
1105
+ Specifies the resources needed when executing this step.
1106
+
1107
+ Use `@resources` to specify the resource requirements
1108
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
1078
1109
 
1110
+ You can choose the compute layer on the command line by executing e.g.
1111
+ ```
1112
+ python myflow.py run --with batch
1113
+ ```
1114
+ or
1115
+ ```
1116
+ python myflow.py run --with kubernetes
1117
+ ```
1118
+ which executes the flow on the desired system using the
1119
+ requirements specified in `@resources`.
1079
1120
 
1080
1121
  Parameters
1081
1122
  ----------
1082
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
1083
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
1084
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
1085
- - `current.checkpoint`
1086
- - `current.model`
1087
- - `current.huggingface_hub`
1088
-
1089
- 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
1090
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
1091
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
1092
-
1093
- temp_dir_root : str, default: None
1094
- The root directory under which `current.model.loaded` will store loaded models
1095
-
1096
-
1097
-
1098
- """
1099
- ...
1100
-
1101
- @typing.overload
1102
- def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1103
- ...
1104
-
1105
- @typing.overload
1106
- def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1107
- ...
1108
-
1109
- 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):
1110
- """
1111
- Enables loading / saving of models within a step.
1112
-
1113
-
1114
- Parameters
1115
- ----------
1116
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
1117
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
1118
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
1119
- - `current.checkpoint`
1120
- - `current.model`
1121
- - `current.huggingface_hub`
1122
-
1123
- 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
1124
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
1125
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
1126
-
1127
- temp_dir_root : str, default: None
1128
- The root directory under which `current.model.loaded` will store loaded models
1129
-
1130
-
1131
-
1123
+ cpu : int, default 1
1124
+ Number of CPUs required for this step.
1125
+ gpu : int, default 0
1126
+ Number of GPUs required for this step.
1127
+ disk : int, optional, default None
1128
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
1129
+ memory : int, default 4096
1130
+ Memory size (in MB) required for this step.
1131
+ shared_memory : int, optional, default None
1132
+ The value for the size (in MiB) of the /dev/shm volume for this step.
1133
+ This parameter maps to the `--shm-size` option in Docker.
1132
1134
  """
1133
1135
  ...
1134
1136
 
1135
1137
  @typing.overload
1136
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1137
- """
1138
- Internal decorator to support Fast bakery
1139
- """
1138
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1140
1139
  ...
1141
1140
 
1142
1141
  @typing.overload
1143
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1142
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1144
1143
  ...
1145
1144
 
1146
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1145
+ def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, cpu: int = 1, gpu: int = 0, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None):
1147
1146
  """
1148
- Internal decorator to support Fast bakery
1147
+ Specifies the resources needed when executing this step.
1148
+
1149
+ Use `@resources` to specify the resource requirements
1150
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
1151
+
1152
+ You can choose the compute layer on the command line by executing e.g.
1153
+ ```
1154
+ python myflow.py run --with batch
1155
+ ```
1156
+ or
1157
+ ```
1158
+ python myflow.py run --with kubernetes
1159
+ ```
1160
+ which executes the flow on the desired system using the
1161
+ requirements specified in `@resources`.
1162
+
1163
+ Parameters
1164
+ ----------
1165
+ cpu : int, default 1
1166
+ Number of CPUs required for this step.
1167
+ gpu : int, default 0
1168
+ Number of GPUs required for this step.
1169
+ disk : int, optional, default None
1170
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
1171
+ memory : int, default 4096
1172
+ Memory size (in MB) required for this step.
1173
+ shared_memory : int, optional, default None
1174
+ The value for the size (in MiB) of the /dev/shm volume for this step.
1175
+ This parameter maps to the `--shm-size` option in Docker.
1149
1176
  """
1150
1177
  ...
1151
1178
 
1152
- @typing.overload
1153
- 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]]]:
1179
+ 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]]]:
1154
1180
  """
1155
- Specifies the Conda environment for the step.
1156
-
1157
- Information in this decorator will augment any
1158
- attributes set in the `@conda_base` flow-level decorator. Hence,
1159
- you can use `@conda_base` to set packages required by all
1160
- steps and use `@conda` to specify step-specific overrides.
1181
+ Decorator that helps cache, version and store models/datasets from huggingface hub.
1161
1182
 
1162
1183
  Parameters
1163
1184
  ----------
1164
- packages : Dict[str, str], default {}
1165
- Packages to use for this step. The key is the name of the package
1166
- and the value is the version to use.
1167
- libraries : Dict[str, str], default {}
1168
- Supported for backward compatibility. When used with packages, packages will take precedence.
1169
- python : str, optional, default None
1170
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1171
- that the version used will correspond to the version of the Python interpreter used to start the run.
1172
- disabled : bool, default False
1173
- If set to True, disables @conda.
1185
+ temp_dir_root : str, optional
1186
+ The root directory that will hold the temporary directory where objects will be downloaded.
1187
+
1188
+
1174
1189
  """
1175
1190
  ...
1176
1191
 
1177
1192
  @typing.overload
1178
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1193
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1194
+ """
1195
+ Internal decorator to support Fast bakery
1196
+ """
1179
1197
  ...
1180
1198
 
1181
1199
  @typing.overload
1182
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1200
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1183
1201
  ...
1184
1202
 
1185
- 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):
1203
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1186
1204
  """
1187
- Specifies the Conda environment for the step.
1188
-
1189
- Information in this decorator will augment any
1190
- attributes set in the `@conda_base` flow-level decorator. Hence,
1191
- you can use `@conda_base` to set packages required by all
1192
- steps and use `@conda` to specify step-specific overrides.
1193
-
1194
- Parameters
1195
- ----------
1196
- packages : Dict[str, str], default {}
1197
- Packages to use for this step. The key is the name of the package
1198
- and the value is the version to use.
1199
- libraries : Dict[str, str], default {}
1200
- Supported for backward compatibility. When used with packages, packages will take precedence.
1201
- python : str, optional, default None
1202
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1203
- that the version used will correspond to the version of the Python interpreter used to start the run.
1204
- disabled : bool, default False
1205
- If set to True, disables @conda.
1205
+ Internal decorator to support Fast bakery
1206
1206
  """
1207
1207
  ...
1208
1208
 
@@ -1309,55 +1309,6 @@ def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *
1309
1309
  """
1310
1310
  ...
1311
1311
 
1312
- @typing.overload
1313
- 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]]:
1314
- """
1315
- Specifies the Conda environment for all steps of the flow.
1316
-
1317
- Use `@conda_base` to set common libraries required by all
1318
- steps and use `@conda` to specify step-specific additions.
1319
-
1320
- Parameters
1321
- ----------
1322
- packages : Dict[str, str], default {}
1323
- Packages to use for this flow. The key is the name of the package
1324
- and the value is the version to use.
1325
- libraries : Dict[str, str], default {}
1326
- Supported for backward compatibility. When used with packages, packages will take precedence.
1327
- python : str, optional, default None
1328
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1329
- that the version used will correspond to the version of the Python interpreter used to start the run.
1330
- disabled : bool, default False
1331
- If set to True, disables Conda.
1332
- """
1333
- ...
1334
-
1335
- @typing.overload
1336
- def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1337
- ...
1338
-
1339
- 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):
1340
- """
1341
- Specifies the Conda environment for all steps of the flow.
1342
-
1343
- Use `@conda_base` to set common libraries required by all
1344
- steps and use `@conda` to specify step-specific additions.
1345
-
1346
- Parameters
1347
- ----------
1348
- packages : Dict[str, str], default {}
1349
- Packages to use for this flow. The key is the name of the package
1350
- and the value is the version to use.
1351
- libraries : Dict[str, str], default {}
1352
- Supported for backward compatibility. When used with packages, packages will take precedence.
1353
- python : str, optional, default None
1354
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1355
- that the version used will correspond to the version of the Python interpreter used to start the run.
1356
- disabled : bool, default False
1357
- If set to True, disables Conda.
1358
- """
1359
- ...
1360
-
1361
1312
  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]]:
1362
1313
  """
1363
1314
  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)
@@ -1400,136 +1351,6 @@ def airflow_s3_key_sensor(*, timeout: int, poke_interval: int, mode: str, expone
1400
1351
  """
1401
1352
  ...
1402
1353
 
1403
- def project(*, name: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1404
- """
1405
- Specifies what flows belong to the same project.
1406
-
1407
- A project-specific namespace is created for all flows that
1408
- use the same `@project(name)`.
1409
-
1410
- Parameters
1411
- ----------
1412
- name : str
1413
- Project name. Make sure that the name is unique amongst all
1414
- projects that use the same production scheduler. The name may
1415
- contain only lowercase alphanumeric characters and underscores.
1416
-
1417
-
1418
- """
1419
- ...
1420
-
1421
- 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]]:
1422
- """
1423
- 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.
1424
- 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.
1425
-
1426
- Parameters
1427
- ----------
1428
- timeout : int
1429
- Time, in seconds before the task times out and fails. (Default: 3600)
1430
- poke_interval : int
1431
- Time in seconds that the job should wait in between each try. (Default: 60)
1432
- mode : str
1433
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1434
- exponential_backoff : bool
1435
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1436
- pool : str
1437
- the slot pool this task should run in,
1438
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1439
- soft_fail : bool
1440
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1441
- name : str
1442
- Name of the sensor on Airflow
1443
- description : str
1444
- Description of sensor in the Airflow UI
1445
- external_dag_id : str
1446
- The dag_id that contains the task you want to wait for.
1447
- external_task_ids : List[str]
1448
- The list of task_ids that you want to wait for.
1449
- If None (default value) the sensor waits for the DAG. (Default: None)
1450
- allowed_states : List[str]
1451
- Iterable of allowed states, (Default: ['success'])
1452
- failed_states : List[str]
1453
- Iterable of failed or dis-allowed states. (Default: None)
1454
- execution_delta : datetime.timedelta
1455
- time difference with the previous execution to look at,
1456
- the default is the same logical date as the current task or DAG. (Default: None)
1457
- check_existence: bool
1458
- Set to True to check if the external task exists or check if
1459
- the DAG to wait for exists. (Default: True)
1460
- """
1461
- ...
1462
-
1463
- @typing.overload
1464
- def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1465
- """
1466
- Specifies the PyPI packages for all steps of the flow.
1467
-
1468
- Use `@pypi_base` to set common packages required by all
1469
- steps and use `@pypi` to specify step-specific overrides.
1470
- Parameters
1471
- ----------
1472
- packages : Dict[str, str], default: {}
1473
- Packages to use for this flow. The key is the name of the package
1474
- and the value is the version to use.
1475
- python : str, optional, default: None
1476
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1477
- that the version used will correspond to the version of the Python interpreter used to start the run.
1478
- """
1479
- ...
1480
-
1481
- @typing.overload
1482
- def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1483
- ...
1484
-
1485
- def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1486
- """
1487
- Specifies the PyPI packages for all steps of the flow.
1488
-
1489
- Use `@pypi_base` to set common packages required by all
1490
- steps and use `@pypi` to specify step-specific overrides.
1491
- Parameters
1492
- ----------
1493
- packages : Dict[str, str], default: {}
1494
- Packages to use for this flow. The key is the name of the package
1495
- and the value is the version to use.
1496
- python : str, optional, default: None
1497
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1498
- that the version used will correspond to the version of the Python interpreter used to start the run.
1499
- """
1500
- ...
1501
-
1502
- def nim(*, models: "list[NIM]", backend: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1503
- """
1504
- This decorator is used to run NIM containers in Metaflow tasks as sidecars.
1505
-
1506
- User code call
1507
- -----------
1508
- @nim(
1509
- models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
1510
- backend='managed'
1511
- )
1512
-
1513
- Valid backend options
1514
- ---------------------
1515
- - 'managed': Outerbounds selects a compute provider based on the model.
1516
- - 🚧 'dataplane': Run in your account.
1517
-
1518
- Valid model options
1519
- ----------------
1520
- - 'meta/llama3-8b-instruct': 8B parameter model
1521
- - 'meta/llama3-70b-instruct': 70B parameter model
1522
- - Upon request, any model here: https://nvcf.ngc.nvidia.com/functions?filter=nvidia-functions
1523
-
1524
- Parameters
1525
- ----------
1526
- models: list[NIM]
1527
- List of NIM containers running models in sidecars.
1528
- backend: str
1529
- Compute provider to run the NIM container.
1530
- """
1531
- ...
1532
-
1533
1354
  @typing.overload
1534
1355
  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]]:
1535
1356
  """
@@ -1674,6 +1495,185 @@ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: t
1674
1495
  """
1675
1496
  ...
1676
1497
 
1498
+ def nim(*, models: "list[NIM]", backend: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1499
+ """
1500
+ This decorator is used to run NIM containers in Metaflow tasks as sidecars.
1501
+
1502
+ User code call
1503
+ -----------
1504
+ @nim(
1505
+ models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
1506
+ backend='managed'
1507
+ )
1508
+
1509
+ Valid backend options
1510
+ ---------------------
1511
+ - 'managed': Outerbounds selects a compute provider based on the model.
1512
+ - 🚧 'dataplane': Run in your account.
1513
+
1514
+ Valid model options
1515
+ ----------------
1516
+ - 'meta/llama3-8b-instruct': 8B parameter model
1517
+ - 'meta/llama3-70b-instruct': 70B parameter model
1518
+ - Upon request, any model here: https://nvcf.ngc.nvidia.com/functions?filter=nvidia-functions
1519
+
1520
+ Parameters
1521
+ ----------
1522
+ models: list[NIM]
1523
+ List of NIM containers running models in sidecars.
1524
+ backend: str
1525
+ Compute provider to run the NIM container.
1526
+ """
1527
+ ...
1528
+
1529
+ @typing.overload
1530
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1531
+ """
1532
+ Specifies the PyPI packages for all steps of the flow.
1533
+
1534
+ Use `@pypi_base` to set common packages required by all
1535
+ steps and use `@pypi` to specify step-specific overrides.
1536
+ Parameters
1537
+ ----------
1538
+ packages : Dict[str, str], default: {}
1539
+ Packages to use for this flow. The key is the name of the package
1540
+ and the value is the version to use.
1541
+ python : str, optional, default: None
1542
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1543
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1544
+ """
1545
+ ...
1546
+
1547
+ @typing.overload
1548
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1549
+ ...
1550
+
1551
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1552
+ """
1553
+ Specifies the PyPI packages for all steps of the flow.
1554
+
1555
+ Use `@pypi_base` to set common packages required by all
1556
+ steps and use `@pypi` to specify step-specific overrides.
1557
+ Parameters
1558
+ ----------
1559
+ packages : Dict[str, str], default: {}
1560
+ Packages to use for this flow. The key is the name of the package
1561
+ and the value is the version to use.
1562
+ python : str, optional, default: None
1563
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1564
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1565
+ """
1566
+ ...
1567
+
1568
+ 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]]:
1569
+ """
1570
+ 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.
1571
+ 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.
1572
+
1573
+ Parameters
1574
+ ----------
1575
+ timeout : int
1576
+ Time, in seconds before the task times out and fails. (Default: 3600)
1577
+ poke_interval : int
1578
+ Time in seconds that the job should wait in between each try. (Default: 60)
1579
+ mode : str
1580
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1581
+ exponential_backoff : bool
1582
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1583
+ pool : str
1584
+ the slot pool this task should run in,
1585
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1586
+ soft_fail : bool
1587
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1588
+ name : str
1589
+ Name of the sensor on Airflow
1590
+ description : str
1591
+ Description of sensor in the Airflow UI
1592
+ external_dag_id : str
1593
+ The dag_id that contains the task you want to wait for.
1594
+ external_task_ids : List[str]
1595
+ The list of task_ids that you want to wait for.
1596
+ If None (default value) the sensor waits for the DAG. (Default: None)
1597
+ allowed_states : List[str]
1598
+ Iterable of allowed states, (Default: ['success'])
1599
+ failed_states : List[str]
1600
+ Iterable of failed or dis-allowed states. (Default: None)
1601
+ execution_delta : datetime.timedelta
1602
+ time difference with the previous execution to look at,
1603
+ the default is the same logical date as the current task or DAG. (Default: None)
1604
+ check_existence: bool
1605
+ Set to True to check if the external task exists or check if
1606
+ the DAG to wait for exists. (Default: True)
1607
+ """
1608
+ ...
1609
+
1610
+ @typing.overload
1611
+ 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]]:
1612
+ """
1613
+ Specifies the Conda environment for all steps of the flow.
1614
+
1615
+ Use `@conda_base` to set common libraries required by all
1616
+ steps and use `@conda` to specify step-specific additions.
1617
+
1618
+ Parameters
1619
+ ----------
1620
+ packages : Dict[str, str], default {}
1621
+ Packages to use for this flow. The key is the name of the package
1622
+ and the value is the version to use.
1623
+ libraries : Dict[str, str], default {}
1624
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1625
+ python : str, optional, default None
1626
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1627
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1628
+ disabled : bool, default False
1629
+ If set to True, disables Conda.
1630
+ """
1631
+ ...
1632
+
1633
+ @typing.overload
1634
+ def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1635
+ ...
1636
+
1637
+ 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):
1638
+ """
1639
+ Specifies the Conda environment for all steps of the flow.
1640
+
1641
+ Use `@conda_base` to set common libraries required by all
1642
+ steps and use `@conda` to specify step-specific additions.
1643
+
1644
+ Parameters
1645
+ ----------
1646
+ packages : Dict[str, str], default {}
1647
+ Packages to use for this flow. The key is the name of the package
1648
+ and the value is the version to use.
1649
+ libraries : Dict[str, str], default {}
1650
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1651
+ python : str, optional, default None
1652
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1653
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1654
+ disabled : bool, default False
1655
+ If set to True, disables Conda.
1656
+ """
1657
+ ...
1658
+
1659
+ def project(*, name: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1660
+ """
1661
+ Specifies what flows belong to the same project.
1662
+
1663
+ A project-specific namespace is created for all flows that
1664
+ use the same `@project(name)`.
1665
+
1666
+ Parameters
1667
+ ----------
1668
+ name : str
1669
+ Project name. Make sure that the name is unique amongst all
1670
+ projects that use the same production scheduler. The name may
1671
+ contain only lowercase alphanumeric characters and underscores.
1672
+
1673
+
1674
+ """
1675
+ ...
1676
+
1677
1677
  def namespace(ns: typing.Optional[str]) -> typing.Optional[str]:
1678
1678
  """
1679
1679
  Switch namespace to the one provided.