ob-metaflow-stubs 6.0.3.113__py2.py3-none-any.whl → 6.0.3.114__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 +676 -673
  2. metaflow-stubs/cards.pyi +5 -5
  3. metaflow-stubs/cli.pyi +2 -2
  4. metaflow-stubs/client/__init__.pyi +3 -3
  5. metaflow-stubs/client/core.pyi +6 -6
  6. metaflow-stubs/client/filecache.pyi +2 -2
  7. metaflow-stubs/clone_util.pyi +2 -2
  8. metaflow-stubs/events.pyi +2 -2
  9. metaflow-stubs/exception.pyi +2 -2
  10. metaflow-stubs/flowspec.pyi +5 -5
  11. metaflow-stubs/generated_for.txt +1 -1
  12. metaflow-stubs/includefile.pyi +4 -4
  13. metaflow-stubs/info_file.pyi +2 -2
  14. metaflow-stubs/metadata/metadata.pyi +3 -3
  15. metaflow-stubs/metadata/util.pyi +2 -2
  16. metaflow-stubs/metaflow_config.pyi +2 -2
  17. metaflow-stubs/metaflow_current.pyi +106 -105
  18. metaflow-stubs/mflog/mflog.pyi +2 -2
  19. metaflow-stubs/multicore_utils.pyi +2 -2
  20. metaflow-stubs/parameters.pyi +4 -4
  21. metaflow-stubs/plugins/__init__.pyi +2 -2
  22. metaflow-stubs/plugins/airflow/__init__.pyi +2 -2
  23. metaflow-stubs/plugins/airflow/airflow_utils.pyi +2 -2
  24. metaflow-stubs/plugins/airflow/exception.pyi +2 -2
  25. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +2 -2
  26. metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi +2 -2
  27. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +3 -3
  28. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +3 -3
  29. metaflow-stubs/plugins/argo/__init__.pyi +2 -2
  30. metaflow-stubs/plugins/argo/argo_client.pyi +11 -2
  31. metaflow-stubs/plugins/argo/argo_events.pyi +2 -2
  32. metaflow-stubs/plugins/argo/argo_workflows.pyi +3 -3
  33. metaflow-stubs/plugins/argo/argo_workflows_cli.pyi +8 -5
  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 +3 -3
  45. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +2 -2
  46. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +2 -2
  47. metaflow-stubs/plugins/aws/step_functions/production_token.pyi +2 -2
  48. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +2 -2
  49. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +2 -2
  50. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +2 -2
  51. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer.pyi +2 -2
  52. metaflow-stubs/plugins/azure/__init__.pyi +2 -2
  53. metaflow-stubs/plugins/azure/azure_credential.pyi +2 -2
  54. metaflow-stubs/plugins/azure/azure_exceptions.pyi +2 -2
  55. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +3 -3
  56. metaflow-stubs/plugins/azure/azure_utils.pyi +2 -2
  57. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +2 -2
  58. metaflow-stubs/plugins/azure/includefile_support.pyi +2 -2
  59. metaflow-stubs/plugins/cards/__init__.pyi +2 -2
  60. metaflow-stubs/plugins/cards/card_cli.pyi +4 -4
  61. metaflow-stubs/plugins/cards/card_client.pyi +3 -3
  62. metaflow-stubs/plugins/cards/card_creator.pyi +2 -2
  63. metaflow-stubs/plugins/cards/card_datastore.pyi +2 -2
  64. metaflow-stubs/plugins/cards/card_decorator.pyi +3 -3
  65. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +2 -2
  66. metaflow-stubs/plugins/cards/card_modules/basic.pyi +3 -3
  67. metaflow-stubs/plugins/cards/card_modules/card.pyi +2 -2
  68. metaflow-stubs/plugins/cards/card_modules/components.pyi +3 -3
  69. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +2 -2
  70. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +2 -2
  71. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +3 -3
  72. metaflow-stubs/plugins/cards/card_resolver.pyi +2 -2
  73. metaflow-stubs/plugins/cards/component_serializer.pyi +4 -4
  74. metaflow-stubs/plugins/cards/exception.pyi +2 -2
  75. metaflow-stubs/plugins/catch_decorator.pyi +3 -3
  76. metaflow-stubs/plugins/datatools/__init__.pyi +2 -2
  77. metaflow-stubs/plugins/datatools/local.pyi +2 -2
  78. metaflow-stubs/plugins/datatools/s3/__init__.pyi +2 -2
  79. metaflow-stubs/plugins/datatools/s3/s3.pyi +6 -6
  80. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +2 -2
  81. metaflow-stubs/plugins/datatools/s3/s3util.pyi +2 -2
  82. metaflow-stubs/plugins/debug_logger.pyi +2 -2
  83. metaflow-stubs/plugins/debug_monitor.pyi +2 -2
  84. metaflow-stubs/plugins/environment_decorator.pyi +2 -2
  85. metaflow-stubs/plugins/events_decorator.pyi +3 -3
  86. metaflow-stubs/plugins/frameworks/__init__.pyi +2 -2
  87. metaflow-stubs/plugins/frameworks/pytorch.pyi +8 -7
  88. metaflow-stubs/plugins/gcp/__init__.pyi +2 -2
  89. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +3 -3
  90. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +2 -2
  91. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +2 -2
  92. metaflow-stubs/plugins/gcp/gs_utils.pyi +2 -2
  93. metaflow-stubs/plugins/gcp/includefile_support.pyi +2 -2
  94. metaflow-stubs/plugins/kubernetes/__init__.pyi +2 -2
  95. metaflow-stubs/plugins/kubernetes/kube_utils.pyi +2 -2
  96. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +2 -2
  97. metaflow-stubs/plugins/kubernetes/kubernetes_cli.pyi +5 -2
  98. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +2 -2
  99. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +6 -3
  100. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +2 -2
  101. metaflow-stubs/plugins/logs_cli.pyi +3 -3
  102. metaflow-stubs/plugins/package_cli.pyi +2 -2
  103. metaflow-stubs/plugins/parallel_decorator.pyi +8 -7
  104. metaflow-stubs/plugins/perimeters.pyi +2 -2
  105. metaflow-stubs/plugins/project_decorator.pyi +3 -3
  106. metaflow-stubs/plugins/pypi/__init__.pyi +2 -2
  107. metaflow-stubs/plugins/pypi/conda_decorator.pyi +2 -2
  108. metaflow-stubs/plugins/pypi/conda_environment.pyi +3 -3
  109. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +2 -2
  110. metaflow-stubs/plugins/pypi/pypi_environment.pyi +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 +4 -4
  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 +2 -2
  122. metaflow-stubs/profilers/__init__.pyi +2 -2
  123. metaflow-stubs/pylint_wrapper.pyi +2 -2
  124. metaflow-stubs/runner/__init__.pyi +2 -2
  125. metaflow-stubs/runner/deployer.pyi +3 -3
  126. metaflow-stubs/runner/metaflow_runner.pyi +5 -8
  127. metaflow-stubs/runner/nbdeploy.pyi +5 -7
  128. metaflow-stubs/runner/nbrun.pyi +5 -7
  129. metaflow-stubs/runner/subprocess_manager.pyi +2 -2
  130. metaflow-stubs/runner/utils.pyi +2 -6
  131. metaflow-stubs/system/__init__.pyi +3 -3
  132. metaflow-stubs/system/system_logger.pyi +2 -2
  133. metaflow-stubs/system/system_monitor.pyi +2 -2
  134. metaflow-stubs/tagging_util.pyi +2 -2
  135. metaflow-stubs/tuple_util.pyi +2 -2
  136. {ob_metaflow_stubs-6.0.3.113.dist-info → ob_metaflow_stubs-6.0.3.114.dist-info}/METADATA +1 -1
  137. ob_metaflow_stubs-6.0.3.114.dist-info/RECORD +140 -0
  138. ob_metaflow_stubs-6.0.3.113.dist-info/RECORD +0 -140
  139. {ob_metaflow_stubs-6.0.3.113.dist-info → ob_metaflow_stubs-6.0.3.114.dist-info}/WHEEL +0 -0
  140. {ob_metaflow_stubs-6.0.3.113.dist-info → ob_metaflow_stubs-6.0.3.114.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.1);ob(v1) #
4
- # Generated on 2024-10-22T22:37:22.374436 #
3
+ # MF version: 2.12.26.1+obcheckpoint(0.1.1);ob(v1) #
4
+ # Generated on 2024-10-23T20:57:01.122821 #
5
5
  ##################################################################################
6
6
 
7
7
  from __future__ import annotations
8
8
 
9
9
  import typing
10
10
  if typing.TYPE_CHECKING:
11
- import datetime
11
+ import metaflow.flowspec
12
12
  import metaflow_extensions.obcheckpoint.plugins.machine_learning_utilities.datastructures
13
13
  import metaflow.datastore.inputs
14
14
  import metaflow.parameters
15
- import metaflow.flowspec
16
- import typing
17
- import metaflow._vendor.click.types
15
+ import metaflow.events
16
+ import metaflow.runner.metaflow_runner
18
17
  import metaflow.client.core
19
18
  import metaflow.metaflow_current
20
- import metaflow.runner.metaflow_runner
21
- import metaflow.events
19
+ import datetime
20
+ import metaflow._vendor.click.types
21
+ import typing
22
22
  FlowSpecDerived = typing.TypeVar("FlowSpecDerived", bound="FlowSpec", contravariant=False, covariant=False)
23
23
  StepFlag = typing.NewType("StepFlag", bool)
24
24
 
@@ -513,440 +513,492 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
513
513
  ...
514
514
 
515
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
+ 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]]]:
517
517
  """
518
- Specifies that the step will success under all circumstances.
518
+ Enables checkpointing for a step.
519
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.
524
520
 
525
521
  Parameters
526
522
  ----------
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.
523
+ load_policy : str, default: "fresh"
524
+ The policy for loading the checkpoint. The following policies are supported:
525
+ - "eager": Loads the the latest available checkpoint within the namespace.
526
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
527
+ will be loaded at the start of the task.
528
+ - "none": Do not load any checkpoint
529
+ - "fresh": Loads the lastest checkpoint created within the running Task.
530
+ This mode helps loading checkpoints across various retry attempts of the same task.
531
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
532
+ created within the task will be loaded when the task is retries execution on failure.
533
+
534
+ temp_dir_root : str, default: None
535
+ The root directory under which `current.checkpoint.directory` will be created.
536
+
537
+
538
+
533
539
  """
534
540
  ...
535
541
 
536
542
  @typing.overload
537
- def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
543
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
538
544
  ...
539
545
 
540
546
  @typing.overload
541
- def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
547
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
542
548
  ...
543
549
 
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):
550
+ 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):
545
551
  """
546
- Specifies that the step will success under all circumstances.
552
+ Enables checkpointing for a step.
547
553
 
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.
552
554
 
553
555
  Parameters
554
556
  ----------
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.
557
+ load_policy : str, default: "fresh"
558
+ The policy for loading the checkpoint. The following policies are supported:
559
+ - "eager": Loads the the latest available checkpoint within the namespace.
560
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
561
+ will be loaded at the start of the task.
562
+ - "none": Do not load any checkpoint
563
+ - "fresh": Loads the lastest checkpoint created within the running Task.
564
+ This mode helps loading checkpoints across various retry attempts of the same task.
565
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
566
+ created within the task will be loaded when the task is retries execution on failure.
567
+
568
+ temp_dir_root : str, default: None
569
+ The root directory under which `current.checkpoint.directory` will be created.
570
+
571
+
572
+
561
573
  """
562
574
  ...
563
575
 
564
576
  @typing.overload
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]]]:
577
+ 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]]]:
566
578
  """
567
- Creates a human-readable report, a Metaflow Card, after this step completes.
568
-
569
- Note that you may add multiple `@card` decorators in a step with different parameters.
579
+ Specifies secrets to be retrieved and injected as environment variables prior to
580
+ the execution of a step.
570
581
 
571
582
  Parameters
572
583
  ----------
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
-
584
+ sources : List[Union[str, Dict[str, Any]]], default: []
585
+ List of secret specs, defining how the secrets are to be retrieved
583
586
  """
584
587
  ...
585
588
 
586
589
  @typing.overload
587
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
590
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
588
591
  ...
589
592
 
590
593
  @typing.overload
591
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
594
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
592
595
  ...
593
596
 
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):
597
+ 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]]] = []):
595
598
  """
596
- Creates a human-readable report, a Metaflow Card, after this step completes.
597
-
598
- Note that you may add multiple `@card` decorators in a step with different parameters.
599
+ Specifies secrets to be retrieved and injected as environment variables prior to
600
+ the execution of a step.
599
601
 
600
602
  Parameters
601
603
  ----------
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
-
604
+ sources : List[Union[str, Dict[str, Any]]], default: []
605
+ List of secret specs, defining how the secrets are to be retrieved
612
606
  """
613
607
  ...
614
608
 
615
- @typing.overload
616
- 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]]]:
609
+ 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]]]:
617
610
  """
618
- Specifies the number of times the task corresponding
619
- to a step needs to be retried.
620
-
621
- This decorator is useful for handling transient errors, such as networking issues.
622
- If your task contains operations that can't be retried safely, e.g. database updates,
623
- it is advisable to annotate it with `@retry(times=0)`.
624
-
625
- This can be used in conjunction with the `@catch` decorator. The `@catch`
626
- decorator will execute a no-op task after all retries have been exhausted,
627
- ensuring that the flow execution can continue.
611
+ Specifies that this step should execute on Kubernetes.
628
612
 
629
613
  Parameters
630
614
  ----------
631
- times : int, default 3
632
- Number of times to retry this task.
633
- minutes_between_retries : int, default 2
634
- Number of minutes between retries.
615
+ cpu : int, default 1
616
+ Number of CPUs required for this step. If `@resources` is
617
+ also present, the maximum value from all decorators is used.
618
+ memory : int, default 4096
619
+ Memory size (in MB) required for this step. If
620
+ `@resources` is also present, the maximum value from all decorators is
621
+ used.
622
+ disk : int, default 10240
623
+ Disk size (in MB) required for this step. If
624
+ `@resources` is also present, the maximum value from all decorators is
625
+ used.
626
+ image : str, optional, default None
627
+ Docker image to use when launching on Kubernetes. If not specified, and
628
+ METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
629
+ not, a default Docker image mapping to the current version of Python is used.
630
+ image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
631
+ If given, the imagePullPolicy to be applied to the Docker image of the step.
632
+ service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
633
+ Kubernetes service account to use when launching pod in Kubernetes.
634
+ secrets : List[str], optional, default None
635
+ Kubernetes secrets to use when launching pod in Kubernetes. These
636
+ secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
637
+ in Metaflow configuration.
638
+ node_selector: Union[Dict[str,str], str], optional, default None
639
+ Kubernetes node selector(s) to apply to the pod running the task.
640
+ Can be passed in as a comma separated string of values e.g. "kubernetes.io/os=linux,kubernetes.io/arch=amd64"
641
+ or as a dictionary {"kubernetes.io/os": "linux", "kubernetes.io/arch": "amd64"}
642
+ namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
643
+ Kubernetes namespace to use when launching pod in Kubernetes.
644
+ gpu : int, optional, default None
645
+ Number of GPUs required for this step. A value of zero implies that
646
+ the scheduled node should not have GPUs.
647
+ gpu_vendor : str, default KUBERNETES_GPU_VENDOR
648
+ The vendor of the GPUs to be used for this step.
649
+ tolerations : List[str], default []
650
+ The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
651
+ Kubernetes tolerations to use when launching pod in Kubernetes.
652
+ use_tmpfs : bool, default False
653
+ This enables an explicit tmpfs mount for this step.
654
+ tmpfs_tempdir : bool, default True
655
+ sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
656
+ tmpfs_size : int, optional, default: None
657
+ The value for the size (in MiB) of the tmpfs mount for this step.
658
+ This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
659
+ memory allocated for this step.
660
+ tmpfs_path : str, optional, default /metaflow_temp
661
+ Path to tmpfs mount for this step.
662
+ persistent_volume_claims : Dict[str, str], optional, default None
663
+ A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
664
+ volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
665
+ shared_memory: int, optional
666
+ Shared memory size (in MiB) required for this step
667
+ port: int, optional
668
+ Port number to specify in the Kubernetes job object
669
+ compute_pool : str, optional, default None
670
+ Compute pool to be used for for this step.
671
+ If not specified, any accessible compute pool within the perimeter is used.
672
+ hostname_resolution_timeout: int, default 10 * 60
673
+ Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
674
+ Only applicable when @parallel is used.
635
675
  """
636
676
  ...
637
677
 
638
678
  @typing.overload
639
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
640
- ...
641
-
642
- @typing.overload
643
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
644
- ...
645
-
646
- 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):
679
+ 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]]]:
647
680
  """
648
- Specifies the number of times the task corresponding
649
- to a step needs to be retried.
650
-
651
- This decorator is useful for handling transient errors, such as networking issues.
652
- If your task contains operations that can't be retried safely, e.g. database updates,
653
- it is advisable to annotate it with `@retry(times=0)`.
681
+ Specifies the Conda environment for the step.
654
682
 
655
- This can be used in conjunction with the `@catch` decorator. The `@catch`
656
- decorator will execute a no-op task after all retries have been exhausted,
657
- ensuring that the flow execution can continue.
683
+ Information in this decorator will augment any
684
+ attributes set in the `@conda_base` flow-level decorator. Hence,
685
+ you can use `@conda_base` to set packages required by all
686
+ steps and use `@conda` to specify step-specific overrides.
658
687
 
659
688
  Parameters
660
689
  ----------
661
- times : int, default 3
662
- Number of times to retry this task.
663
- minutes_between_retries : int, default 2
664
- Number of minutes between retries.
690
+ packages : Dict[str, str], default {}
691
+ Packages to use for this step. The key is the name of the package
692
+ and the value is the version to use.
693
+ libraries : Dict[str, str], default {}
694
+ Supported for backward compatibility. When used with packages, packages will take precedence.
695
+ python : str, optional, default None
696
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
697
+ that the version used will correspond to the version of the Python interpreter used to start the run.
698
+ disabled : bool, default False
699
+ If set to True, disables @conda.
665
700
  """
666
701
  ...
667
702
 
668
703
  @typing.overload
669
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
670
- """
671
- Decorator prototype for all step decorators. This function gets specialized
672
- and imported for all decorators types by _import_plugin_decorators().
673
- """
704
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
674
705
  ...
675
706
 
676
707
  @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
- """
708
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
685
709
  ...
686
710
 
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]]]:
711
+ 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):
689
712
  """
690
- Enables loading / saving of models within a step.
713
+ Specifies the Conda environment for the step.
691
714
 
715
+ Information in this decorator will augment any
716
+ attributes set in the `@conda_base` flow-level decorator. Hence,
717
+ you can use `@conda_base` to set packages required by all
718
+ steps and use `@conda` to specify step-specific overrides.
692
719
 
693
720
  Parameters
694
721
  ----------
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.
722
+ packages : Dict[str, str], default {}
723
+ Packages to use for this step. The key is the name of the package
724
+ and the value is the version to use.
725
+ libraries : Dict[str, str], default {}
726
+ Supported for backward compatibility. When used with packages, packages will take precedence.
727
+ python : str, optional, default None
728
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
729
+ that the version used will correspond to the version of the Python interpreter used to start the run.
730
+ disabled : bool, default False
731
+ If set to True, disables @conda.
732
+ """
733
+ ...
734
+
735
+ @typing.overload
736
+ 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]]]:
737
+ """
738
+ Specifies a timeout for your step.
705
739
 
706
- temp_dir_root : str, default: None
707
- The root directory under which `current.model.loaded` will store loaded models
740
+ This decorator is useful if this step may hang indefinitely.
708
741
 
742
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
743
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
744
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
709
745
 
746
+ Note that all the values specified in parameters are added together so if you specify
747
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
710
748
 
749
+ Parameters
750
+ ----------
751
+ seconds : int, default 0
752
+ Number of seconds to wait prior to timing out.
753
+ minutes : int, default 0
754
+ Number of minutes to wait prior to timing out.
755
+ hours : int, default 0
756
+ Number of hours to wait prior to timing out.
711
757
  """
712
758
  ...
713
759
 
714
760
  @typing.overload
715
- def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
761
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
716
762
  ...
717
763
 
718
764
  @typing.overload
719
- def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
765
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
720
766
  ...
721
767
 
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):
768
+ 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):
723
769
  """
724
- Enables loading / saving of models within a step.
725
-
726
-
727
- Parameters
728
- ----------
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.
770
+ Specifies a timeout for your step.
739
771
 
740
- temp_dir_root : str, default: None
741
- The root directory under which `current.model.loaded` will store loaded models
772
+ This decorator is useful if this step may hang indefinitely.
742
773
 
774
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
775
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
776
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
743
777
 
778
+ Note that all the values specified in parameters are added together so if you specify
779
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
744
780
 
781
+ Parameters
782
+ ----------
783
+ seconds : int, default 0
784
+ Number of seconds to wait prior to timing out.
785
+ minutes : int, default 0
786
+ Number of minutes to wait prior to timing out.
787
+ hours : int, default 0
788
+ Number of hours to wait prior to timing out.
745
789
  """
746
790
  ...
747
791
 
748
792
  @typing.overload
749
- 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]]]:
793
+ 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]]]:
750
794
  """
751
- Specifies the PyPI packages for the step.
795
+ Specifies the resources needed when executing this step.
752
796
 
753
- Information in this decorator will augment any
754
- attributes set in the `@pyi_base` flow-level decorator. Hence,
755
- you can use `@pypi_base` to set packages required by all
756
- steps and use `@pypi` to specify step-specific overrides.
797
+ Use `@resources` to specify the resource requirements
798
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
799
+
800
+ You can choose the compute layer on the command line by executing e.g.
801
+ ```
802
+ python myflow.py run --with batch
803
+ ```
804
+ or
805
+ ```
806
+ python myflow.py run --with kubernetes
807
+ ```
808
+ which executes the flow on the desired system using the
809
+ requirements specified in `@resources`.
757
810
 
758
811
  Parameters
759
812
  ----------
760
- packages : Dict[str, str], default: {}
761
- Packages to use for this step. The key is the name of the package
762
- and the value is the version to use.
763
- python : str, optional, default: None
764
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
765
- that the version used will correspond to the version of the Python interpreter used to start the run.
813
+ cpu : int, default 1
814
+ Number of CPUs required for this step.
815
+ gpu : int, default 0
816
+ Number of GPUs required for this step.
817
+ disk : int, optional, default None
818
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
819
+ memory : int, default 4096
820
+ Memory size (in MB) required for this step.
821
+ shared_memory : int, optional, default None
822
+ The value for the size (in MiB) of the /dev/shm volume for this step.
823
+ This parameter maps to the `--shm-size` option in Docker.
766
824
  """
767
825
  ...
768
826
 
769
827
  @typing.overload
770
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
828
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
771
829
  ...
772
830
 
773
831
  @typing.overload
774
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
832
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
775
833
  ...
776
834
 
777
- 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):
835
+ 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):
778
836
  """
779
- Specifies the PyPI packages for the step.
837
+ Specifies the resources needed when executing this step.
780
838
 
781
- Information in this decorator will augment any
782
- attributes set in the `@pyi_base` flow-level decorator. Hence,
783
- you can use `@pypi_base` to set packages required by all
784
- steps and use `@pypi` to specify step-specific overrides.
839
+ Use `@resources` to specify the resource requirements
840
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
841
+
842
+ You can choose the compute layer on the command line by executing e.g.
843
+ ```
844
+ python myflow.py run --with batch
845
+ ```
846
+ or
847
+ ```
848
+ python myflow.py run --with kubernetes
849
+ ```
850
+ which executes the flow on the desired system using the
851
+ requirements specified in `@resources`.
785
852
 
786
853
  Parameters
787
854
  ----------
788
- packages : Dict[str, str], default: {}
789
- Packages to use for this step. The key is the name of the package
790
- and the value is the version to use.
791
- python : str, optional, default: None
792
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
793
- that the version used will correspond to the version of the Python interpreter used to start the run.
855
+ cpu : int, default 1
856
+ Number of CPUs required for this step.
857
+ gpu : int, default 0
858
+ Number of GPUs required for this step.
859
+ disk : int, optional, default None
860
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
861
+ memory : int, default 4096
862
+ Memory size (in MB) required for this step.
863
+ shared_memory : int, optional, default None
864
+ The value for the size (in MiB) of the /dev/shm volume for this step.
865
+ This parameter maps to the `--shm-size` option in Docker.
794
866
  """
795
867
  ...
796
868
 
797
869
  @typing.overload
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]]]:
870
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
799
871
  """
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
872
+ Internal decorator to support Fast bakery
807
873
  """
808
874
  ...
809
875
 
810
876
  @typing.overload
811
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
812
- ...
813
-
814
- @typing.overload
815
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
877
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
816
878
  ...
817
879
 
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]]] = []):
880
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
819
881
  """
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
882
+ Internal decorator to support Fast bakery
827
883
  """
828
884
  ...
829
885
 
830
886
  @typing.overload
831
- 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]]]:
887
+ 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]]]:
832
888
  """
833
- Specifies a timeout for your step.
834
-
835
- This decorator is useful if this step may hang indefinitely.
836
-
837
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
838
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
839
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
889
+ Specifies that the step will success under all circumstances.
840
890
 
841
- Note that all the values specified in parameters are added together so if you specify
842
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
891
+ The decorator will create an optional artifact, specified by `var`, which
892
+ contains the exception raised. You can use it to detect the presence
893
+ of errors, indicating that all happy-path artifacts produced by the step
894
+ are missing.
843
895
 
844
896
  Parameters
845
897
  ----------
846
- seconds : int, default 0
847
- Number of seconds to wait prior to timing out.
848
- minutes : int, default 0
849
- Number of minutes to wait prior to timing out.
850
- hours : int, default 0
851
- Number of hours to wait prior to timing out.
898
+ var : str, optional, default None
899
+ Name of the artifact in which to store the caught exception.
900
+ If not specified, the exception is not stored.
901
+ print_exception : bool, default True
902
+ Determines whether or not the exception is printed to
903
+ stdout when caught.
852
904
  """
853
905
  ...
854
906
 
855
907
  @typing.overload
856
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
908
+ def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
857
909
  ...
858
910
 
859
911
  @typing.overload
860
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
912
+ def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
861
913
  ...
862
914
 
863
- 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):
915
+ 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):
864
916
  """
865
- Specifies a timeout for your step.
866
-
867
- This decorator is useful if this step may hang indefinitely.
868
-
869
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
870
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
871
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
917
+ Specifies that the step will success under all circumstances.
872
918
 
873
- Note that all the values specified in parameters are added together so if you specify
874
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
919
+ The decorator will create an optional artifact, specified by `var`, which
920
+ contains the exception raised. You can use it to detect the presence
921
+ of errors, indicating that all happy-path artifacts produced by the step
922
+ are missing.
875
923
 
876
924
  Parameters
877
925
  ----------
878
- seconds : int, default 0
879
- Number of seconds to wait prior to timing out.
880
- minutes : int, default 0
881
- Number of minutes to wait prior to timing out.
882
- hours : int, default 0
883
- Number of hours to wait prior to timing out.
926
+ var : str, optional, default None
927
+ Name of the artifact in which to store the caught exception.
928
+ If not specified, the exception is not stored.
929
+ print_exception : bool, default True
930
+ Determines whether or not the exception is printed to
931
+ stdout when caught.
884
932
  """
885
933
  ...
886
934
 
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]]]:
935
+ @typing.overload
936
+ 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]]]:
888
937
  """
889
- Specifies that this step should execute on Kubernetes.
938
+ Creates a human-readable report, a Metaflow Card, after this step completes.
939
+
940
+ Note that you may add multiple `@card` decorators in a step with different parameters.
890
941
 
891
942
  Parameters
892
- ----------
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.
943
+ ----------
944
+ type : str, default 'default'
945
+ Card type.
946
+ id : str, optional, default None
947
+ If multiple cards are present, use this id to identify this card.
948
+ options : Dict[str, Any], default {}
949
+ Options passed to the card. The contents depend on the card type.
950
+ timeout : int, default 45
951
+ Interrupt reporting if it takes more than this many seconds.
952
+
953
+
954
+ """
955
+ ...
956
+
957
+ @typing.overload
958
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
959
+ ...
960
+
961
+ @typing.overload
962
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
963
+ ...
964
+
965
+ 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):
966
+ """
967
+ Creates a human-readable report, a Metaflow Card, after this step completes.
968
+
969
+ Note that you may add multiple `@card` decorators in a step with different parameters.
970
+
971
+ Parameters
972
+ ----------
973
+ type : str, default 'default'
974
+ Card type.
975
+ id : str, optional, default None
976
+ If multiple cards are present, use this id to identify this card.
977
+ options : Dict[str, Any], default {}
978
+ Options passed to the card. The contents depend on the card type.
979
+ timeout : int, default 45
980
+ Interrupt reporting if it takes more than this many seconds.
981
+
982
+
983
+ """
984
+ ...
985
+
986
+ @typing.overload
987
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
988
+ """
989
+ Decorator prototype for all step decorators. This function gets specialized
990
+ and imported for all decorators types by _import_plugin_decorators().
991
+ """
992
+ ...
993
+
994
+ @typing.overload
995
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
996
+ ...
997
+
998
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
999
+ """
1000
+ Decorator prototype for all step decorators. This function gets specialized
1001
+ and imported for all decorators types by _import_plugin_decorators().
950
1002
  """
951
1003
  ...
952
1004
 
@@ -982,227 +1034,178 @@ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], Non
982
1034
  ...
983
1035
 
984
1036
  @typing.overload
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]]]:
1037
+ 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]]]:
986
1038
  """
987
- Specifies the Conda environment for the step.
1039
+ Specifies the PyPI packages for the step.
988
1040
 
989
1041
  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.
1042
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
1043
+ you can use `@pypi_base` to set packages required by all
1044
+ steps and use `@pypi` to specify step-specific overrides.
993
1045
 
994
1046
  Parameters
995
1047
  ----------
996
- packages : Dict[str, str], default {}
1048
+ packages : Dict[str, str], default: {}
997
1049
  Packages to use for this step. The key is the name of the package
998
1050
  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
1051
+ python : str, optional, default: None
1002
1052
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
1003
1053
  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.
1006
1054
  """
1007
1055
  ...
1008
1056
 
1009
1057
  @typing.overload
1010
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1058
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1011
1059
  ...
1012
1060
 
1013
1061
  @typing.overload
1014
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1062
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1015
1063
  ...
1016
1064
 
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):
1065
+ 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):
1018
1066
  """
1019
- Specifies the Conda environment for the step.
1067
+ Specifies the PyPI packages for the step.
1020
1068
 
1021
1069
  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.
1070
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
1071
+ you can use `@pypi_base` to set packages required by all
1072
+ steps and use `@pypi` to specify step-specific overrides.
1025
1073
 
1026
1074
  Parameters
1027
1075
  ----------
1028
- packages : Dict[str, str], default {}
1076
+ packages : Dict[str, str], default: {}
1029
1077
  Packages to use for this step. The key is the name of the package
1030
1078
  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
1079
+ python : str, optional, default: None
1034
1080
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
1035
1081
  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.
1038
- """
1039
- ...
1040
-
1041
- @typing.overload
1042
- 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]]]:
1043
- """
1044
- Enables checkpointing for a step.
1045
-
1046
-
1047
- Parameters
1048
- ----------
1049
- load_policy : str, default: "fresh"
1050
- The policy for loading the checkpoint. The following policies are supported:
1051
- - "eager": Loads the the latest available checkpoint within the namespace.
1052
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
1053
- will be loaded at the start of the task.
1054
- - "none": Do not load any checkpoint
1055
- - "fresh": Loads the lastest checkpoint created within the running Task.
1056
- This mode helps loading checkpoints across various retry attempts of the same task.
1057
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
1058
- created within the task will be loaded when the task is retries execution on failure.
1059
-
1060
- temp_dir_root : str, default: None
1061
- The root directory under which `current.checkpoint.directory` will be created.
1062
-
1063
-
1064
-
1065
1082
  """
1066
1083
  ...
1067
1084
 
1068
- @typing.overload
1069
- def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1070
- ...
1071
-
1072
- @typing.overload
1073
- def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1074
- ...
1075
-
1076
- 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):
1085
+ 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]]]:
1077
1086
  """
1078
- Enables checkpointing for a step.
1079
-
1087
+ Decorator that helps cache, version and store models/datasets from huggingface hub.
1080
1088
 
1081
1089
  Parameters
1082
1090
  ----------
1083
- load_policy : str, default: "fresh"
1084
- The policy for loading the checkpoint. The following policies are supported:
1085
- - "eager": Loads the the latest available checkpoint within the namespace.
1086
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
1087
- will be loaded at the start of the task.
1088
- - "none": Do not load any checkpoint
1089
- - "fresh": Loads the lastest checkpoint created within the running Task.
1090
- This mode helps loading checkpoints across various retry attempts of the same task.
1091
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
1092
- created within the task will be loaded when the task is retries execution on failure.
1093
-
1094
- temp_dir_root : str, default: None
1095
- The root directory under which `current.checkpoint.directory` will be created.
1096
-
1091
+ temp_dir_root : str, optional
1092
+ The root directory that will hold the temporary directory where objects will be downloaded.
1097
1093
 
1098
1094
 
1099
1095
  """
1100
1096
  ...
1101
1097
 
1102
1098
  @typing.overload
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]]]:
1099
+ 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]]]:
1104
1100
  """
1105
- Specifies the resources needed when executing this step.
1101
+ Specifies the number of times the task corresponding
1102
+ to a step needs to be retried.
1106
1103
 
1107
- Use `@resources` to specify the resource requirements
1108
- independently of the specific compute layer (`@batch`, `@kubernetes`).
1104
+ This decorator is useful for handling transient errors, such as networking issues.
1105
+ If your task contains operations that can't be retried safely, e.g. database updates,
1106
+ it is advisable to annotate it with `@retry(times=0)`.
1109
1107
 
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`.
1108
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
1109
+ decorator will execute a no-op task after all retries have been exhausted,
1110
+ ensuring that the flow execution can continue.
1120
1111
 
1121
1112
  Parameters
1122
1113
  ----------
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.
1114
+ times : int, default 3
1115
+ Number of times to retry this task.
1116
+ minutes_between_retries : int, default 2
1117
+ Number of minutes between retries.
1134
1118
  """
1135
1119
  ...
1136
1120
 
1137
1121
  @typing.overload
1138
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1122
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1139
1123
  ...
1140
1124
 
1141
1125
  @typing.overload
1142
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1126
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1143
1127
  ...
1144
1128
 
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):
1129
+ 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):
1146
1130
  """
1147
- Specifies the resources needed when executing this step.
1131
+ Specifies the number of times the task corresponding
1132
+ to a step needs to be retried.
1148
1133
 
1149
- Use `@resources` to specify the resource requirements
1150
- independently of the specific compute layer (`@batch`, `@kubernetes`).
1134
+ This decorator is useful for handling transient errors, such as networking issues.
1135
+ If your task contains operations that can't be retried safely, e.g. database updates,
1136
+ it is advisable to annotate it with `@retry(times=0)`.
1151
1137
 
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`.
1138
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
1139
+ decorator will execute a no-op task after all retries have been exhausted,
1140
+ ensuring that the flow execution can continue.
1162
1141
 
1163
1142
  Parameters
1164
1143
  ----------
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.
1144
+ times : int, default 3
1145
+ Number of times to retry this task.
1146
+ minutes_between_retries : int, default 2
1147
+ Number of minutes between retries.
1176
1148
  """
1177
1149
  ...
1178
1150
 
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]]]:
1151
+ @typing.overload
1152
+ 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]]]:
1180
1153
  """
1181
- Decorator that helps cache, version and store models/datasets from huggingface hub.
1154
+ Enables loading / saving of models within a step.
1155
+
1156
+
1157
+ Parameters
1158
+ ----------
1159
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
1160
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
1161
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
1162
+ - `current.checkpoint`
1163
+ - `current.model`
1164
+ - `current.huggingface_hub`
1165
+
1166
+ 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
1167
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
1168
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
1169
+
1170
+ temp_dir_root : str, default: None
1171
+ The root directory under which `current.model.loaded` will store loaded models
1182
1172
 
1183
- Parameters
1184
- ----------
1185
- temp_dir_root : str, optional
1186
- The root directory that will hold the temporary directory where objects will be downloaded.
1187
1173
 
1188
1174
 
1189
1175
  """
1190
1176
  ...
1191
1177
 
1192
1178
  @typing.overload
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
+ def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1197
1180
  ...
1198
1181
 
1199
1182
  @typing.overload
1200
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1183
+ def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1201
1184
  ...
1202
1185
 
1203
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1186
+ 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):
1204
1187
  """
1205
- Internal decorator to support Fast bakery
1188
+ Enables loading / saving of models within a step.
1189
+
1190
+
1191
+ Parameters
1192
+ ----------
1193
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
1194
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
1195
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
1196
+ - `current.checkpoint`
1197
+ - `current.model`
1198
+ - `current.huggingface_hub`
1199
+
1200
+ 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
1201
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
1202
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
1203
+
1204
+ temp_dir_root : str, default: None
1205
+ The root directory under which `current.model.loaded` will store loaded models
1206
+
1207
+
1208
+
1206
1209
  """
1207
1210
  ...
1208
1211
 
@@ -1298,14 +1301,242 @@ def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *
1298
1301
 
1299
1302
  Parameters
1300
1303
  ----------
1301
- flow : Union[str, Dict[str, str]], optional, default None
1302
- Upstream flow dependency for this flow.
1303
- flows : List[Union[str, Dict[str, str]]], default []
1304
- Upstream flow dependencies for this flow.
1305
- options : Dict[str, Any], default {}
1306
- Backend-specific configuration for tuning eventing behavior.
1307
-
1308
-
1304
+ flow : Union[str, Dict[str, str]], optional, default None
1305
+ Upstream flow dependency for this flow.
1306
+ flows : List[Union[str, Dict[str, str]]], default []
1307
+ Upstream flow dependencies for this flow.
1308
+ options : Dict[str, Any], default {}
1309
+ Backend-specific configuration for tuning eventing behavior.
1310
+
1311
+
1312
+ """
1313
+ ...
1314
+
1315
+ def nim(*, models: "list[NIM]", backend: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1316
+ """
1317
+ This decorator is used to run NIM containers in Metaflow tasks as sidecars.
1318
+
1319
+ User code call
1320
+ -----------
1321
+ @nim(
1322
+ models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
1323
+ backend='managed'
1324
+ )
1325
+
1326
+ Valid backend options
1327
+ ---------------------
1328
+ - 'managed': Outerbounds selects a compute provider based on the model.
1329
+ - 🚧 'dataplane': Run in your account.
1330
+
1331
+ Valid model options
1332
+ ----------------
1333
+ - 'meta/llama3-8b-instruct': 8B parameter model
1334
+ - 'meta/llama3-70b-instruct': 70B parameter model
1335
+ - Upon request, any model here: https://nvcf.ngc.nvidia.com/functions?filter=nvidia-functions
1336
+
1337
+ Parameters
1338
+ ----------
1339
+ models: list[NIM]
1340
+ List of NIM containers running models in sidecars.
1341
+ backend: str
1342
+ Compute provider to run the NIM container.
1343
+ """
1344
+ ...
1345
+
1346
+ def project(*, name: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1347
+ """
1348
+ Specifies what flows belong to the same project.
1349
+
1350
+ A project-specific namespace is created for all flows that
1351
+ use the same `@project(name)`.
1352
+
1353
+ Parameters
1354
+ ----------
1355
+ name : str
1356
+ Project name. Make sure that the name is unique amongst all
1357
+ projects that use the same production scheduler. The name may
1358
+ contain only lowercase alphanumeric characters and underscores.
1359
+
1360
+
1361
+ """
1362
+ ...
1363
+
1364
+ 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]]:
1365
+ """
1366
+ 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.
1367
+ 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.
1368
+
1369
+ Parameters
1370
+ ----------
1371
+ timeout : int
1372
+ Time, in seconds before the task times out and fails. (Default: 3600)
1373
+ poke_interval : int
1374
+ Time in seconds that the job should wait in between each try. (Default: 60)
1375
+ mode : str
1376
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1377
+ exponential_backoff : bool
1378
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1379
+ pool : str
1380
+ the slot pool this task should run in,
1381
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1382
+ soft_fail : bool
1383
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1384
+ name : str
1385
+ Name of the sensor on Airflow
1386
+ description : str
1387
+ Description of sensor in the Airflow UI
1388
+ external_dag_id : str
1389
+ The dag_id that contains the task you want to wait for.
1390
+ external_task_ids : List[str]
1391
+ The list of task_ids that you want to wait for.
1392
+ If None (default value) the sensor waits for the DAG. (Default: None)
1393
+ allowed_states : List[str]
1394
+ Iterable of allowed states, (Default: ['success'])
1395
+ failed_states : List[str]
1396
+ Iterable of failed or dis-allowed states. (Default: None)
1397
+ execution_delta : datetime.timedelta
1398
+ time difference with the previous execution to look at,
1399
+ the default is the same logical date as the current task or DAG. (Default: None)
1400
+ check_existence: bool
1401
+ Set to True to check if the external task exists or check if
1402
+ the DAG to wait for exists. (Default: True)
1403
+ """
1404
+ ...
1405
+
1406
+ @typing.overload
1407
+ 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]]:
1408
+ """
1409
+ Specifies the Conda environment for all steps of the flow.
1410
+
1411
+ Use `@conda_base` to set common libraries required by all
1412
+ steps and use `@conda` to specify step-specific additions.
1413
+
1414
+ Parameters
1415
+ ----------
1416
+ packages : Dict[str, str], default {}
1417
+ Packages to use for this flow. The key is the name of the package
1418
+ and the value is the version to use.
1419
+ libraries : Dict[str, str], default {}
1420
+ Supported for backward compatibility. When used with packages, packages will take precedence.
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
+ disabled : bool, default False
1425
+ If set to True, disables Conda.
1426
+ """
1427
+ ...
1428
+
1429
+ @typing.overload
1430
+ def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1431
+ ...
1432
+
1433
+ 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):
1434
+ """
1435
+ Specifies the Conda environment for all steps of the flow.
1436
+
1437
+ Use `@conda_base` to set common libraries required by all
1438
+ steps and use `@conda` to specify step-specific additions.
1439
+
1440
+ Parameters
1441
+ ----------
1442
+ packages : Dict[str, str], default {}
1443
+ Packages to use for this flow. The key is the name of the package
1444
+ and the value is the version to use.
1445
+ libraries : Dict[str, str], default {}
1446
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1447
+ python : str, optional, default None
1448
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1449
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1450
+ disabled : bool, default False
1451
+ If set to True, disables Conda.
1452
+ """
1453
+ ...
1454
+
1455
+ @typing.overload
1456
+ 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]]:
1457
+ """
1458
+ Specifies the times when the flow should be run when running on a
1459
+ production scheduler.
1460
+
1461
+ Parameters
1462
+ ----------
1463
+ hourly : bool, default False
1464
+ Run the workflow hourly.
1465
+ daily : bool, default True
1466
+ Run the workflow daily.
1467
+ weekly : bool, default False
1468
+ Run the workflow weekly.
1469
+ cron : str, optional, default None
1470
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1471
+ specified by this expression.
1472
+ timezone : str, optional, default None
1473
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1474
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1475
+ """
1476
+ ...
1477
+
1478
+ @typing.overload
1479
+ def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1480
+ ...
1481
+
1482
+ def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None):
1483
+ """
1484
+ Specifies the times when the flow should be run when running on a
1485
+ production scheduler.
1486
+
1487
+ Parameters
1488
+ ----------
1489
+ hourly : bool, default False
1490
+ Run the workflow hourly.
1491
+ daily : bool, default True
1492
+ Run the workflow daily.
1493
+ weekly : bool, default False
1494
+ Run the workflow weekly.
1495
+ cron : str, optional, default None
1496
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1497
+ specified by this expression.
1498
+ timezone : str, optional, default None
1499
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1500
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1501
+ """
1502
+ ...
1503
+
1504
+ @typing.overload
1505
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1506
+ """
1507
+ Specifies the PyPI packages for all steps of the flow.
1508
+
1509
+ Use `@pypi_base` to set common packages required by all
1510
+ steps and use `@pypi` to specify step-specific overrides.
1511
+ Parameters
1512
+ ----------
1513
+ packages : Dict[str, str], default: {}
1514
+ Packages to use for this flow. The key is the name of the package
1515
+ and the value is the version to use.
1516
+ python : str, optional, default: None
1517
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1518
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1519
+ """
1520
+ ...
1521
+
1522
+ @typing.overload
1523
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1524
+ ...
1525
+
1526
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1527
+ """
1528
+ Specifies the PyPI packages for all steps of the flow.
1529
+
1530
+ Use `@pypi_base` to set common packages required by all
1531
+ steps and use `@pypi` to specify step-specific overrides.
1532
+ Parameters
1533
+ ----------
1534
+ packages : Dict[str, str], default: {}
1535
+ Packages to use for this flow. The key is the name of the package
1536
+ and the value is the version to use.
1537
+ python : str, optional, default: None
1538
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1539
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1309
1540
  """
1310
1541
  ...
1311
1542
 
@@ -1351,55 +1582,6 @@ def airflow_s3_key_sensor(*, timeout: int, poke_interval: int, mode: str, expone
1351
1582
  """
1352
1583
  ...
1353
1584
 
1354
- @typing.overload
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]]:
1356
- """
1357
- Specifies the times when the flow should be run when running on a
1358
- production scheduler.
1359
-
1360
- Parameters
1361
- ----------
1362
- hourly : bool, default False
1363
- Run the workflow hourly.
1364
- daily : bool, default True
1365
- Run the workflow daily.
1366
- weekly : bool, default False
1367
- Run the workflow weekly.
1368
- cron : str, optional, default None
1369
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1370
- specified by this expression.
1371
- timezone : str, optional, default None
1372
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1373
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1374
- """
1375
- ...
1376
-
1377
- @typing.overload
1378
- def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1379
- ...
1380
-
1381
- def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None):
1382
- """
1383
- Specifies the times when the flow should be run when running on a
1384
- production scheduler.
1385
-
1386
- Parameters
1387
- ----------
1388
- hourly : bool, default False
1389
- Run the workflow hourly.
1390
- daily : bool, default True
1391
- Run the workflow daily.
1392
- weekly : bool, default False
1393
- Run the workflow weekly.
1394
- cron : str, optional, default None
1395
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1396
- specified by this expression.
1397
- timezone : str, optional, default None
1398
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1399
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1400
- """
1401
- ...
1402
-
1403
1585
  @typing.overload
1404
1586
  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]]:
1405
1587
  """
@@ -1495,185 +1677,6 @@ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: t
1495
1677
  """
1496
1678
  ...
1497
1679
 
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
1680
  def namespace(ns: typing.Optional[str]) -> typing.Optional[str]:
1678
1681
  """
1679
1682
  Switch namespace to the one provided.
@@ -2908,15 +2911,15 @@ class NBRunner(object, metaclass=type):
2908
2911
  Additional environment variables to set for the Run. This overrides the
2909
2912
  environment set for this process.
2910
2913
  base_dir : Optional[str], default None
2911
- The directory to run the subprocess in; if not specified, a temporary
2912
- directory is used.
2914
+ The directory to run the subprocess in; if not specified, the current
2915
+ working directory is used.
2913
2916
  file_read_timeout : int, default 3600
2914
2917
  The timeout until which we try to read the runner attribute file.
2915
2918
  **kwargs : Any
2916
2919
  Additional arguments that you would pass to `python myflow.py` before
2917
2920
  the `run` command.
2918
2921
  """
2919
- def __init__(self, flow, show_output: bool = True, profile: typing.Optional[str] = None, env: typing.Optional[typing.Dict] = None, base_dir: str = "/tmp", file_read_timeout: int = 3600, **kwargs):
2922
+ def __init__(self, flow, show_output: bool = True, profile: typing.Optional[str] = None, env: typing.Optional[typing.Dict] = None, base_dir: typing.Optional[str] = None, file_read_timeout: int = 3600, **kwargs):
2920
2923
  ...
2921
2924
  def nbrun(self, **kwargs):
2922
2925
  """
@@ -3097,13 +3100,13 @@ class NBDeployer(object, metaclass=type):
3097
3100
  Additional environment variables to set. This overrides the
3098
3101
  environment set for this process.
3099
3102
  base_dir : Optional[str], default None
3100
- The directory to run the subprocess in; if not specified, a temporary
3101
- directory is used.
3103
+ The directory to run the subprocess in; if not specified, the current
3104
+ working directory is used.
3102
3105
  **kwargs : Any
3103
3106
  Additional arguments that you would pass to `python myflow.py` i.e. options
3104
3107
  listed in `python myflow.py --help`
3105
3108
  """
3106
- def __init__(self, flow, show_output: bool = True, profile: typing.Optional[str] = None, env: typing.Optional[typing.Dict] = None, base_dir: str = "/tmp", file_read_timeout: int = 3600, **kwargs):
3109
+ def __init__(self, flow, show_output: bool = True, profile: typing.Optional[str] = None, env: typing.Optional[typing.Dict] = None, base_dir: typing.Optional[str] = None, file_read_timeout: int = 3600, **kwargs):
3107
3110
  ...
3108
3111
  def cleanup(self):
3109
3112
  """