ob-metaflow-stubs 6.0.3.103rc3__py2.py3-none-any.whl → 6.0.3.104__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 +497 -497
  2. metaflow-stubs/cards.pyi +4 -4
  3. metaflow-stubs/cli.pyi +2 -2
  4. metaflow-stubs/client/__init__.pyi +3 -3
  5. metaflow-stubs/client/core.pyi +5 -5
  6. metaflow-stubs/client/filecache.pyi +2 -2
  7. metaflow-stubs/clone_util.pyi +1 -1
  8. metaflow-stubs/events.pyi +1 -1
  9. metaflow-stubs/exception.pyi +1 -1
  10. metaflow-stubs/flowspec.pyi +4 -4
  11. metaflow-stubs/generated_for.txt +1 -1
  12. metaflow-stubs/includefile.pyi +3 -3
  13. metaflow-stubs/info_file.pyi +1 -1
  14. metaflow-stubs/metadata/metadata.pyi +2 -2
  15. metaflow-stubs/metadata/util.pyi +1 -1
  16. metaflow-stubs/metaflow_config.pyi +1 -1
  17. metaflow-stubs/metaflow_current.pyi +86 -86
  18. metaflow-stubs/mflog/mflog.pyi +1 -1
  19. metaflow-stubs/multicore_utils.pyi +1 -1
  20. metaflow-stubs/parameters.pyi +3 -3
  21. metaflow-stubs/plugins/__init__.pyi +2 -2
  22. metaflow-stubs/plugins/airflow/__init__.pyi +1 -1
  23. metaflow-stubs/plugins/airflow/airflow_utils.pyi +1 -1
  24. metaflow-stubs/plugins/airflow/exception.pyi +1 -1
  25. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +1 -1
  26. metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi +2 -2
  27. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +2 -2
  28. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +2 -2
  29. metaflow-stubs/plugins/argo/__init__.pyi +1 -1
  30. metaflow-stubs/plugins/argo/argo_client.pyi +1 -1
  31. metaflow-stubs/plugins/argo/argo_events.pyi +1 -1
  32. metaflow-stubs/plugins/argo/argo_workflows.pyi +3 -3
  33. metaflow-stubs/plugins/argo/argo_workflows_cli.pyi +5 -5
  34. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +5 -5
  35. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +2 -2
  36. metaflow-stubs/plugins/aws/__init__.pyi +1 -1
  37. metaflow-stubs/plugins/aws/aws_client.pyi +1 -1
  38. metaflow-stubs/plugins/aws/aws_utils.pyi +1 -1
  39. metaflow-stubs/plugins/aws/batch/__init__.pyi +1 -1
  40. metaflow-stubs/plugins/aws/batch/batch.pyi +1 -1
  41. metaflow-stubs/plugins/aws/batch/batch_client.pyi +1 -1
  42. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +2 -2
  43. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +1 -1
  44. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +1 -1
  45. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +1 -1
  46. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +1 -1
  47. metaflow-stubs/plugins/aws/step_functions/production_token.pyi +1 -1
  48. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +1 -1
  49. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +1 -1
  50. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +1 -1
  51. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer.pyi +2 -2
  52. metaflow-stubs/plugins/azure/__init__.pyi +1 -1
  53. metaflow-stubs/plugins/azure/azure_credential.pyi +1 -1
  54. metaflow-stubs/plugins/azure/azure_exceptions.pyi +1 -1
  55. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +1 -1
  56. metaflow-stubs/plugins/azure/azure_utils.pyi +1 -1
  57. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +1 -1
  58. metaflow-stubs/plugins/azure/includefile_support.pyi +1 -1
  59. metaflow-stubs/plugins/cards/__init__.pyi +1 -1
  60. metaflow-stubs/plugins/cards/card_cli.pyi +4 -4
  61. metaflow-stubs/plugins/cards/card_client.pyi +2 -2
  62. metaflow-stubs/plugins/cards/card_creator.pyi +1 -1
  63. metaflow-stubs/plugins/cards/card_datastore.pyi +1 -1
  64. metaflow-stubs/plugins/cards/card_decorator.pyi +2 -2
  65. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +1 -1
  66. metaflow-stubs/plugins/cards/card_modules/basic.pyi +2 -2
  67. metaflow-stubs/plugins/cards/card_modules/card.pyi +1 -1
  68. metaflow-stubs/plugins/cards/card_modules/components.pyi +3 -3
  69. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +1 -1
  70. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +1 -1
  71. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +1 -1
  72. metaflow-stubs/plugins/cards/card_resolver.pyi +1 -1
  73. metaflow-stubs/plugins/cards/component_serializer.pyi +2 -2
  74. metaflow-stubs/plugins/cards/exception.pyi +1 -1
  75. metaflow-stubs/plugins/catch_decorator.pyi +3 -3
  76. metaflow-stubs/plugins/datatools/__init__.pyi +2 -2
  77. metaflow-stubs/plugins/datatools/local.pyi +1 -1
  78. metaflow-stubs/plugins/datatools/s3/__init__.pyi +2 -2
  79. metaflow-stubs/plugins/datatools/s3/s3.pyi +4 -4
  80. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +1 -1
  81. metaflow-stubs/plugins/datatools/s3/s3util.pyi +1 -1
  82. metaflow-stubs/plugins/debug_logger.pyi +1 -1
  83. metaflow-stubs/plugins/debug_monitor.pyi +1 -1
  84. metaflow-stubs/plugins/environment_decorator.pyi +1 -1
  85. metaflow-stubs/plugins/events_decorator.pyi +2 -2
  86. metaflow-stubs/plugins/frameworks/__init__.pyi +1 -1
  87. metaflow-stubs/plugins/frameworks/pytorch.pyi +3 -3
  88. metaflow-stubs/plugins/gcp/__init__.pyi +1 -1
  89. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +1 -1
  90. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +1 -1
  91. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +1 -1
  92. metaflow-stubs/plugins/gcp/gs_utils.pyi +1 -1
  93. metaflow-stubs/plugins/gcp/includefile_support.pyi +1 -1
  94. metaflow-stubs/plugins/kubernetes/__init__.pyi +1 -1
  95. metaflow-stubs/plugins/kubernetes/kube_utils.pyi +1 -1
  96. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +2 -2
  97. metaflow-stubs/plugins/kubernetes/kubernetes_cli.pyi +3 -3
  98. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +1 -1
  99. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +2 -2
  100. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +1 -1
  101. metaflow-stubs/plugins/logs_cli.pyi +2 -2
  102. metaflow-stubs/plugins/package_cli.pyi +1 -1
  103. metaflow-stubs/plugins/parallel_decorator.pyi +2 -2
  104. metaflow-stubs/plugins/perimeters.pyi +1 -1
  105. metaflow-stubs/plugins/project_decorator.pyi +2 -2
  106. metaflow-stubs/plugins/pypi/__init__.pyi +1 -1
  107. metaflow-stubs/plugins/pypi/conda_decorator.pyi +1 -1
  108. metaflow-stubs/plugins/pypi/conda_environment.pyi +2 -2
  109. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +1 -1
  110. metaflow-stubs/plugins/pypi/pypi_environment.pyi +2 -2
  111. metaflow-stubs/plugins/pypi/utils.pyi +1 -1
  112. metaflow-stubs/plugins/resources_decorator.pyi +1 -1
  113. metaflow-stubs/plugins/retry_decorator.pyi +1 -1
  114. metaflow-stubs/plugins/secrets/__init__.pyi +1 -1
  115. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +1 -1
  116. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +1 -1
  117. metaflow-stubs/plugins/storage_executor.pyi +1 -1
  118. metaflow-stubs/plugins/tag_cli.pyi +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 +1 -1
  122. metaflow-stubs/profilers/__init__.pyi +1 -1
  123. metaflow-stubs/pylint_wrapper.pyi +1 -1
  124. metaflow-stubs/runner/__init__.pyi +1 -1
  125. metaflow-stubs/runner/deployer.pyi +2 -2
  126. metaflow-stubs/runner/metaflow_runner.pyi +4 -4
  127. metaflow-stubs/runner/nbdeploy.pyi +1 -1
  128. metaflow-stubs/runner/nbrun.pyi +1 -1
  129. metaflow-stubs/runner/subprocess_manager.pyi +1 -1
  130. metaflow-stubs/runner/utils.pyi +1 -1
  131. metaflow-stubs/system/__init__.pyi +2 -2
  132. metaflow-stubs/system/system_logger.pyi +1 -1
  133. metaflow-stubs/system/system_monitor.pyi +2 -2
  134. metaflow-stubs/tagging_util.pyi +1 -1
  135. metaflow-stubs/tuple_util.pyi +1 -1
  136. {ob_metaflow_stubs-6.0.3.103rc3.dist-info → ob_metaflow_stubs-6.0.3.104.dist-info}/METADATA +1 -1
  137. ob_metaflow_stubs-6.0.3.104.dist-info/RECORD +140 -0
  138. ob_metaflow_stubs-6.0.3.103rc3.dist-info/RECORD +0 -140
  139. {ob_metaflow_stubs-6.0.3.103rc3.dist-info → ob_metaflow_stubs-6.0.3.104.dist-info}/WHEEL +0 -0
  140. {ob_metaflow_stubs-6.0.3.103rc3.dist-info → ob_metaflow_stubs-6.0.3.104.dist-info}/top_level.txt +0 -0
@@ -1,24 +1,24 @@
1
1
  ##################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
3
  # MF version: 2.12.22.1+obcheckpoint(0.0.11);ob(v1) #
4
- # Generated on 2024-09-20T21:42:44.532164 #
4
+ # Generated on 2024-09-20T21:46:07.326886 #
5
5
  ##################################################################################
6
6
 
7
7
  from __future__ import annotations
8
8
 
9
9
  import typing
10
10
  if typing.TYPE_CHECKING:
11
- import metaflow.flowspec
12
- import metaflow.parameters
13
- import typing
14
- import metaflow.datastore.inputs
11
+ import metaflow.client.core
15
12
  import metaflow._vendor.click.types
16
13
  import metaflow.runner.metaflow_runner
17
- import metaflow_extensions.obcheckpoint.plugins.machine_learning_utilities.datastructures
18
- import datetime
14
+ import metaflow.parameters
19
15
  import metaflow.events
20
16
  import metaflow.metaflow_current
21
- import metaflow.client.core
17
+ import datetime
18
+ import metaflow.datastore.inputs
19
+ import metaflow_extensions.obcheckpoint.plugins.machine_learning_utilities.datastructures
20
+ import metaflow.flowspec
21
+ import typing
22
22
  FlowSpecDerived = typing.TypeVar("FlowSpecDerived", bound="FlowSpec", contravariant=False, covariant=False)
23
23
  StepFlag = typing.NewType("StepFlag", bool)
24
24
 
@@ -434,6 +434,80 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
434
434
  """
435
435
  ...
436
436
 
437
+ @typing.overload
438
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
439
+ """
440
+ Internal decorator to support Fast bakery
441
+ """
442
+ ...
443
+
444
+ @typing.overload
445
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
446
+ ...
447
+
448
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
449
+ """
450
+ Internal decorator to support Fast bakery
451
+ """
452
+ ...
453
+
454
+ @typing.overload
455
+ 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]]]:
456
+ """
457
+ Specifies the Conda environment for the step.
458
+
459
+ Information in this decorator will augment any
460
+ attributes set in the `@conda_base` flow-level decorator. Hence,
461
+ you can use `@conda_base` to set packages required by all
462
+ steps and use `@conda` to specify step-specific overrides.
463
+
464
+ Parameters
465
+ ----------
466
+ packages : Dict[str, str], default {}
467
+ Packages to use for this step. The key is the name of the package
468
+ and the value is the version to use.
469
+ libraries : Dict[str, str], default {}
470
+ Supported for backward compatibility. When used with packages, packages will take precedence.
471
+ python : str, optional, default None
472
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
473
+ that the version used will correspond to the version of the Python interpreter used to start the run.
474
+ disabled : bool, default False
475
+ If set to True, disables @conda.
476
+ """
477
+ ...
478
+
479
+ @typing.overload
480
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
481
+ ...
482
+
483
+ @typing.overload
484
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
485
+ ...
486
+
487
+ 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):
488
+ """
489
+ Specifies the Conda environment for the step.
490
+
491
+ Information in this decorator will augment any
492
+ attributes set in the `@conda_base` flow-level decorator. Hence,
493
+ you can use `@conda_base` to set packages required by all
494
+ steps and use `@conda` to specify step-specific overrides.
495
+
496
+ Parameters
497
+ ----------
498
+ packages : Dict[str, str], default {}
499
+ Packages to use for this step. The key is the name of the package
500
+ and the value is the version to use.
501
+ libraries : Dict[str, str], default {}
502
+ Supported for backward compatibility. When used with packages, packages will take precedence.
503
+ python : str, optional, default None
504
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
505
+ that the version used will correspond to the version of the Python interpreter used to start the run.
506
+ disabled : bool, default False
507
+ If set to True, disables @conda.
508
+ """
509
+ ...
510
+
437
511
  @typing.overload
438
512
  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]]]:
439
513
  """
@@ -484,412 +558,284 @@ def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
484
558
  ...
485
559
 
486
560
  @typing.overload
487
- def environment(*, vars: typing.Dict[str, str] = {}) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
561
+ 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]]]:
488
562
  """
489
- Specifies environment variables to be set prior to the execution of a step.
563
+ Specifies secrets to be retrieved and injected as environment variables prior to
564
+ the execution of a step.
490
565
 
491
566
  Parameters
492
567
  ----------
493
- vars : Dict[str, str], default {}
494
- Dictionary of environment variables to set.
568
+ sources : List[Union[str, Dict[str, Any]]], default: []
569
+ List of secret specs, defining how the secrets are to be retrieved
495
570
  """
496
571
  ...
497
572
 
498
573
  @typing.overload
499
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
574
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
500
575
  ...
501
576
 
502
577
  @typing.overload
503
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
578
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
504
579
  ...
505
580
 
506
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
581
+ 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]]] = []):
507
582
  """
508
- Specifies environment variables to be set prior to the execution of a step.
583
+ Specifies secrets to be retrieved and injected as environment variables prior to
584
+ the execution of a step.
509
585
 
510
586
  Parameters
511
587
  ----------
512
- vars : Dict[str, str], default {}
513
- Dictionary of environment variables to set.
588
+ sources : List[Union[str, Dict[str, Any]]], default: []
589
+ List of secret specs, defining how the secrets are to be retrieved
514
590
  """
515
591
  ...
516
592
 
517
593
  @typing.overload
518
- 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]]]:
594
+ 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]]]:
519
595
  """
520
- Specifies a timeout for your step.
521
-
522
- This decorator is useful if this step may hang indefinitely.
523
-
524
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
525
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
526
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
596
+ Enables loading / saving of models within a step.
527
597
 
528
- Note that all the values specified in parameters are added together so if you specify
529
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
530
598
 
531
599
  Parameters
532
600
  ----------
533
- seconds : int, default 0
534
- Number of seconds to wait prior to timing out.
535
- minutes : int, default 0
536
- Number of minutes to wait prior to timing out.
537
- hours : int, default 0
538
- Number of hours to wait prior to timing out.
601
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
602
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
603
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
604
+ - `current.checkpoint`
605
+ - `current.model`
606
+ - `current.huggingface_hub`
607
+
608
+ 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
609
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
610
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
611
+
612
+ temp_dir_root : str, default: None
613
+ The root directory under which `current.model.loaded` will store loaded models
614
+
615
+
616
+
539
617
  """
540
618
  ...
541
619
 
542
620
  @typing.overload
543
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
621
+ def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
544
622
  ...
545
623
 
546
624
  @typing.overload
547
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
625
+ def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
548
626
  ...
549
627
 
550
- 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):
628
+ 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):
551
629
  """
552
- Specifies a timeout for your step.
553
-
554
- This decorator is useful if this step may hang indefinitely.
555
-
556
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
557
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
558
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
630
+ Enables loading / saving of models within a step.
559
631
 
560
- Note that all the values specified in parameters are added together so if you specify
561
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
562
632
 
563
633
  Parameters
564
634
  ----------
565
- seconds : int, default 0
566
- Number of seconds to wait prior to timing out.
567
- minutes : int, default 0
568
- Number of minutes to wait prior to timing out.
569
- hours : int, default 0
570
- Number of hours to wait prior to timing out.
635
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
636
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
637
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
638
+ - `current.checkpoint`
639
+ - `current.model`
640
+ - `current.huggingface_hub`
641
+
642
+ 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
643
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
644
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
645
+
646
+ temp_dir_root : str, default: None
647
+ The root directory under which `current.model.loaded` will store loaded models
648
+
649
+
650
+
571
651
  """
572
652
  ...
573
653
 
574
654
  @typing.overload
575
- 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]]]:
655
+ 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]]]:
576
656
  """
577
- Specifies the number of times the task corresponding
578
- to a step needs to be retried.
579
-
580
- This decorator is useful for handling transient errors, such as networking issues.
581
- If your task contains operations that can't be retried safely, e.g. database updates,
582
- it is advisable to annotate it with `@retry(times=0)`.
657
+ Creates a human-readable report, a Metaflow Card, after this step completes.
583
658
 
584
- This can be used in conjunction with the `@catch` decorator. The `@catch`
585
- decorator will execute a no-op task after all retries have been exhausted,
586
- ensuring that the flow execution can continue.
659
+ Note that you may add multiple `@card` decorators in a step with different parameters.
587
660
 
588
661
  Parameters
589
662
  ----------
590
- times : int, default 3
591
- Number of times to retry this task.
592
- minutes_between_retries : int, default 2
593
- Number of minutes between retries.
663
+ type : str, default 'default'
664
+ Card type.
665
+ id : str, optional, default None
666
+ If multiple cards are present, use this id to identify this card.
667
+ options : Dict[str, Any], default {}
668
+ Options passed to the card. The contents depend on the card type.
669
+ timeout : int, default 45
670
+ Interrupt reporting if it takes more than this many seconds.
671
+
672
+
594
673
  """
595
674
  ...
596
675
 
597
676
  @typing.overload
598
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
677
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
599
678
  ...
600
679
 
601
680
  @typing.overload
602
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
681
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
603
682
  ...
604
683
 
605
- 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):
684
+ 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):
606
685
  """
607
- Specifies the number of times the task corresponding
608
- to a step needs to be retried.
609
-
610
- This decorator is useful for handling transient errors, such as networking issues.
611
- If your task contains operations that can't be retried safely, e.g. database updates,
612
- it is advisable to annotate it with `@retry(times=0)`.
686
+ Creates a human-readable report, a Metaflow Card, after this step completes.
613
687
 
614
- This can be used in conjunction with the `@catch` decorator. The `@catch`
615
- decorator will execute a no-op task after all retries have been exhausted,
616
- ensuring that the flow execution can continue.
688
+ Note that you may add multiple `@card` decorators in a step with different parameters.
617
689
 
618
690
  Parameters
619
691
  ----------
620
- times : int, default 3
621
- Number of times to retry this task.
622
- minutes_between_retries : int, default 2
623
- Number of minutes between retries.
624
- """
625
- ...
626
-
627
- @typing.overload
628
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
629
- """
630
- Internal decorator to support Fast bakery
692
+ type : str, default 'default'
693
+ Card type.
694
+ id : str, optional, default None
695
+ If multiple cards are present, use this id to identify this card.
696
+ options : Dict[str, Any], default {}
697
+ Options passed to the card. The contents depend on the card type.
698
+ timeout : int, default 45
699
+ Interrupt reporting if it takes more than this many seconds.
700
+
701
+
631
702
  """
632
703
  ...
633
704
 
634
705
  @typing.overload
635
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
636
- ...
637
-
638
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
639
- """
640
- Internal decorator to support Fast bakery
706
+ 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]]]:
641
707
  """
642
- ...
643
-
644
- @typing.overload
645
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
646
- """
647
- Decorator prototype for all step decorators. This function gets specialized
648
- and imported for all decorators types by _import_plugin_decorators().
649
- """
650
- ...
651
-
652
- @typing.overload
653
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
654
- ...
655
-
656
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
657
- """
658
- Decorator prototype for all step decorators. This function gets specialized
659
- and imported for all decorators types by _import_plugin_decorators().
660
- """
661
- ...
662
-
663
- @typing.overload
664
- 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]]]:
665
- """
666
- Specifies the resources needed when executing this step.
667
-
668
- Use `@resources` to specify the resource requirements
669
- independently of the specific compute layer (`@batch`, `@kubernetes`).
708
+ Specifies the PyPI packages for the step.
670
709
 
671
- You can choose the compute layer on the command line by executing e.g.
672
- ```
673
- python myflow.py run --with batch
674
- ```
675
- or
676
- ```
677
- python myflow.py run --with kubernetes
678
- ```
679
- which executes the flow on the desired system using the
680
- requirements specified in `@resources`.
710
+ Information in this decorator will augment any
711
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
712
+ you can use `@pypi_base` to set packages required by all
713
+ steps and use `@pypi` to specify step-specific overrides.
681
714
 
682
715
  Parameters
683
716
  ----------
684
- cpu : int, default 1
685
- Number of CPUs required for this step.
686
- gpu : int, default 0
687
- Number of GPUs required for this step.
688
- disk : int, optional, default None
689
- Disk size (in MB) required for this step. Only applies on Kubernetes.
690
- memory : int, default 4096
691
- Memory size (in MB) required for this step.
692
- shared_memory : int, optional, default None
693
- The value for the size (in MiB) of the /dev/shm volume for this step.
694
- This parameter maps to the `--shm-size` option in Docker.
717
+ packages : Dict[str, str], default: {}
718
+ Packages to use for this step. The key is the name of the package
719
+ and the value is the version to use.
720
+ python : str, optional, default: None
721
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
722
+ that the version used will correspond to the version of the Python interpreter used to start the run.
695
723
  """
696
724
  ...
697
725
 
698
726
  @typing.overload
699
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
727
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
700
728
  ...
701
729
 
702
730
  @typing.overload
703
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
731
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
704
732
  ...
705
733
 
706
- 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):
734
+ 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):
707
735
  """
708
- Specifies the resources needed when executing this step.
709
-
710
- Use `@resources` to specify the resource requirements
711
- independently of the specific compute layer (`@batch`, `@kubernetes`).
736
+ Specifies the PyPI packages for the step.
712
737
 
713
- You can choose the compute layer on the command line by executing e.g.
714
- ```
715
- python myflow.py run --with batch
716
- ```
717
- or
718
- ```
719
- python myflow.py run --with kubernetes
720
- ```
721
- which executes the flow on the desired system using the
722
- requirements specified in `@resources`.
738
+ Information in this decorator will augment any
739
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
740
+ you can use `@pypi_base` to set packages required by all
741
+ steps and use `@pypi` to specify step-specific overrides.
723
742
 
724
743
  Parameters
725
744
  ----------
726
- cpu : int, default 1
727
- Number of CPUs required for this step.
728
- gpu : int, default 0
729
- Number of GPUs required for this step.
730
- disk : int, optional, default None
731
- Disk size (in MB) required for this step. Only applies on Kubernetes.
732
- memory : int, default 4096
733
- Memory size (in MB) required for this step.
734
- shared_memory : int, optional, default None
735
- The value for the size (in MiB) of the /dev/shm volume for this step.
736
- This parameter maps to the `--shm-size` option in Docker.
745
+ packages : Dict[str, str], default: {}
746
+ Packages to use for this step. The key is the name of the package
747
+ and the value is the version to use.
748
+ python : str, optional, default: None
749
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
750
+ that the version used will correspond to the version of the Python interpreter used to start the run.
737
751
  """
738
752
  ...
739
753
 
740
754
  @typing.overload
741
- 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]]]:
755
+ 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]]]:
742
756
  """
743
- Specifies secrets to be retrieved and injected as environment variables prior to
744
- the execution of a step.
757
+ Specifies a timeout for your step.
758
+
759
+ This decorator is useful if this step may hang indefinitely.
760
+
761
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
762
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
763
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
764
+
765
+ Note that all the values specified in parameters are added together so if you specify
766
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
745
767
 
746
768
  Parameters
747
769
  ----------
748
- sources : List[Union[str, Dict[str, Any]]], default: []
749
- List of secret specs, defining how the secrets are to be retrieved
770
+ seconds : int, default 0
771
+ Number of seconds to wait prior to timing out.
772
+ minutes : int, default 0
773
+ Number of minutes to wait prior to timing out.
774
+ hours : int, default 0
775
+ Number of hours to wait prior to timing out.
750
776
  """
751
777
  ...
752
778
 
753
779
  @typing.overload
754
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
780
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
755
781
  ...
756
782
 
757
783
  @typing.overload
758
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
784
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
759
785
  ...
760
786
 
761
- 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]]] = []):
787
+ 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):
762
788
  """
763
- Specifies secrets to be retrieved and injected as environment variables prior to
764
- the execution of a step.
789
+ Specifies a timeout for your step.
765
790
 
766
- Parameters
767
- ----------
768
- sources : List[Union[str, Dict[str, Any]]], default: []
769
- List of secret specs, defining how the secrets are to be retrieved
770
- """
771
- ...
772
-
773
- 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]]]:
774
- """
775
- Specifies that this step should execute on Kubernetes.
791
+ This decorator is useful if this step may hang indefinitely.
792
+
793
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
794
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
795
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
796
+
797
+ Note that all the values specified in parameters are added together so if you specify
798
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
776
799
 
777
800
  Parameters
778
801
  ----------
779
- cpu : int, default 1
780
- Number of CPUs required for this step. If `@resources` is
781
- also present, the maximum value from all decorators is used.
782
- memory : int, default 4096
783
- Memory size (in MB) required for this step. If
784
- `@resources` is also present, the maximum value from all decorators is
785
- used.
786
- disk : int, default 10240
787
- Disk size (in MB) required for this step. If
788
- `@resources` is also present, the maximum value from all decorators is
789
- used.
790
- image : str, optional, default None
791
- Docker image to use when launching on Kubernetes. If not specified, and
792
- METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
793
- not, a default Docker image mapping to the current version of Python is used.
794
- image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
795
- If given, the imagePullPolicy to be applied to the Docker image of the step.
796
- service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
797
- Kubernetes service account to use when launching pod in Kubernetes.
798
- secrets : List[str], optional, default None
799
- Kubernetes secrets to use when launching pod in Kubernetes. These
800
- secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
801
- in Metaflow configuration.
802
- node_selector: Union[Dict[str,str], str], optional, default None
803
- Kubernetes node selector(s) to apply to the pod running the task.
804
- Can be passed in as a comma separated string of values e.g. "kubernetes.io/os=linux,kubernetes.io/arch=amd64"
805
- or as a dictionary {"kubernetes.io/os": "linux", "kubernetes.io/arch": "amd64"}
806
- namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
807
- Kubernetes namespace to use when launching pod in Kubernetes.
808
- gpu : int, optional, default None
809
- Number of GPUs required for this step. A value of zero implies that
810
- the scheduled node should not have GPUs.
811
- gpu_vendor : str, default KUBERNETES_GPU_VENDOR
812
- The vendor of the GPUs to be used for this step.
813
- tolerations : List[str], default []
814
- The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
815
- Kubernetes tolerations to use when launching pod in Kubernetes.
816
- use_tmpfs : bool, default False
817
- This enables an explicit tmpfs mount for this step.
818
- tmpfs_tempdir : bool, default True
819
- sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
820
- tmpfs_size : int, optional, default: None
821
- The value for the size (in MiB) of the tmpfs mount for this step.
822
- This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
823
- memory allocated for this step.
824
- tmpfs_path : str, optional, default /metaflow_temp
825
- Path to tmpfs mount for this step.
826
- persistent_volume_claims : Dict[str, str], optional, default None
827
- A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
828
- volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
829
- shared_memory: int, optional
830
- Shared memory size (in MiB) required for this step
831
- port: int, optional
832
- Port number to specify in the Kubernetes job object
833
- compute_pool : str, optional, default None
834
- Compute pool to be used for for this step.
835
- If not specified, any accessible compute pool within the perimeter is used.
802
+ seconds : int, default 0
803
+ Number of seconds to wait prior to timing out.
804
+ minutes : int, default 0
805
+ Number of minutes to wait prior to timing out.
806
+ hours : int, default 0
807
+ Number of hours to wait prior to timing out.
836
808
  """
837
809
  ...
838
810
 
839
811
  @typing.overload
840
- 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]]]:
812
+ def environment(*, vars: typing.Dict[str, str] = {}) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
841
813
  """
842
- Specifies the Conda environment for the step.
843
-
844
- Information in this decorator will augment any
845
- attributes set in the `@conda_base` flow-level decorator. Hence,
846
- you can use `@conda_base` to set packages required by all
847
- steps and use `@conda` to specify step-specific overrides.
814
+ Specifies environment variables to be set prior to the execution of a step.
848
815
 
849
816
  Parameters
850
817
  ----------
851
- packages : Dict[str, str], default {}
852
- Packages to use for this step. The key is the name of the package
853
- and the value is the version to use.
854
- libraries : Dict[str, str], default {}
855
- Supported for backward compatibility. When used with packages, packages will take precedence.
856
- python : str, optional, default None
857
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
858
- that the version used will correspond to the version of the Python interpreter used to start the run.
859
- disabled : bool, default False
860
- If set to True, disables @conda.
818
+ vars : Dict[str, str], default {}
819
+ Dictionary of environment variables to set.
861
820
  """
862
821
  ...
863
822
 
864
823
  @typing.overload
865
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
824
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
866
825
  ...
867
826
 
868
827
  @typing.overload
869
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
828
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
870
829
  ...
871
830
 
872
- 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):
831
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
873
832
  """
874
- Specifies the Conda environment for the step.
875
-
876
- Information in this decorator will augment any
877
- attributes set in the `@conda_base` flow-level decorator. Hence,
878
- you can use `@conda_base` to set packages required by all
879
- steps and use `@conda` to specify step-specific overrides.
833
+ Specifies environment variables to be set prior to the execution of a step.
880
834
 
881
835
  Parameters
882
836
  ----------
883
- packages : Dict[str, str], default {}
884
- Packages to use for this step. The key is the name of the package
885
- and the value is the version to use.
886
- libraries : Dict[str, str], default {}
887
- Supported for backward compatibility. When used with packages, packages will take precedence.
888
- python : str, optional, default None
889
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
890
- that the version used will correspond to the version of the Python interpreter used to start the run.
891
- disabled : bool, default False
892
- If set to True, disables @conda.
837
+ vars : Dict[str, str], default {}
838
+ Dictionary of environment variables to set.
893
839
  """
894
840
  ...
895
841
 
@@ -955,163 +901,217 @@ def checkpoint(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None
955
901
  ...
956
902
 
957
903
  @typing.overload
958
- 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]]]:
904
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
959
905
  """
960
- Enables loading / saving of models within a step.
961
-
962
-
963
- Parameters
964
- ----------
965
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
966
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
967
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
968
- - `current.checkpoint`
969
- - `current.model`
970
- - `current.huggingface_hub`
971
-
972
- 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
973
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
974
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
975
-
976
- temp_dir_root : str, default: None
977
- The root directory under which `current.model.loaded` will store loaded models
978
-
979
-
980
-
906
+ Decorator prototype for all step decorators. This function gets specialized
907
+ and imported for all decorators types by _import_plugin_decorators().
981
908
  """
982
909
  ...
983
910
 
984
911
  @typing.overload
985
- def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
986
- ...
987
-
988
- @typing.overload
989
- def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
912
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
990
913
  ...
991
914
 
992
- 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):
915
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
993
916
  """
994
- Enables loading / saving of models within a step.
995
-
996
-
997
- Parameters
998
- ----------
999
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
1000
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
1001
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
1002
- - `current.checkpoint`
1003
- - `current.model`
1004
- - `current.huggingface_hub`
1005
-
1006
- 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
1007
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
1008
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
1009
-
1010
- temp_dir_root : str, default: None
1011
- The root directory under which `current.model.loaded` will store loaded models
1012
-
1013
-
1014
-
917
+ Decorator prototype for all step decorators. This function gets specialized
918
+ and imported for all decorators types by _import_plugin_decorators().
1015
919
  """
1016
920
  ...
1017
921
 
1018
922
  @typing.overload
1019
- 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]]]:
923
+ 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]]]:
1020
924
  """
1021
- Creates a human-readable report, a Metaflow Card, after this step completes.
925
+ Specifies the number of times the task corresponding
926
+ to a step needs to be retried.
1022
927
 
1023
- Note that you may add multiple `@card` decorators in a step with different parameters.
928
+ This decorator is useful for handling transient errors, such as networking issues.
929
+ If your task contains operations that can't be retried safely, e.g. database updates,
930
+ it is advisable to annotate it with `@retry(times=0)`.
931
+
932
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
933
+ decorator will execute a no-op task after all retries have been exhausted,
934
+ ensuring that the flow execution can continue.
1024
935
 
1025
936
  Parameters
1026
937
  ----------
1027
- type : str, default 'default'
1028
- Card type.
1029
- id : str, optional, default None
1030
- If multiple cards are present, use this id to identify this card.
1031
- options : Dict[str, Any], default {}
1032
- Options passed to the card. The contents depend on the card type.
1033
- timeout : int, default 45
1034
- Interrupt reporting if it takes more than this many seconds.
1035
-
1036
-
938
+ times : int, default 3
939
+ Number of times to retry this task.
940
+ minutes_between_retries : int, default 2
941
+ Number of minutes between retries.
1037
942
  """
1038
943
  ...
1039
944
 
1040
945
  @typing.overload
1041
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
946
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1042
947
  ...
1043
948
 
1044
949
  @typing.overload
1045
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
950
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1046
951
  ...
1047
952
 
1048
- 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):
953
+ 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):
1049
954
  """
1050
- Creates a human-readable report, a Metaflow Card, after this step completes.
955
+ Specifies the number of times the task corresponding
956
+ to a step needs to be retried.
1051
957
 
1052
- Note that you may add multiple `@card` decorators in a step with different parameters.
958
+ This decorator is useful for handling transient errors, such as networking issues.
959
+ If your task contains operations that can't be retried safely, e.g. database updates,
960
+ it is advisable to annotate it with `@retry(times=0)`.
961
+
962
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
963
+ decorator will execute a no-op task after all retries have been exhausted,
964
+ ensuring that the flow execution can continue.
1053
965
 
1054
966
  Parameters
1055
967
  ----------
1056
- type : str, default 'default'
1057
- Card type.
1058
- id : str, optional, default None
1059
- If multiple cards are present, use this id to identify this card.
1060
- options : Dict[str, Any], default {}
1061
- Options passed to the card. The contents depend on the card type.
1062
- timeout : int, default 45
1063
- Interrupt reporting if it takes more than this many seconds.
1064
-
1065
-
968
+ times : int, default 3
969
+ Number of times to retry this task.
970
+ minutes_between_retries : int, default 2
971
+ Number of minutes between retries.
1066
972
  """
1067
973
  ...
1068
974
 
1069
975
  @typing.overload
1070
- 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]]]:
976
+ 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]]]:
1071
977
  """
1072
- Specifies the PyPI packages for the step.
978
+ Specifies the resources needed when executing this step.
1073
979
 
1074
- Information in this decorator will augment any
1075
- attributes set in the `@pyi_base` flow-level decorator. Hence,
1076
- you can use `@pypi_base` to set packages required by all
1077
- steps and use `@pypi` to specify step-specific overrides.
980
+ Use `@resources` to specify the resource requirements
981
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
982
+
983
+ You can choose the compute layer on the command line by executing e.g.
984
+ ```
985
+ python myflow.py run --with batch
986
+ ```
987
+ or
988
+ ```
989
+ python myflow.py run --with kubernetes
990
+ ```
991
+ which executes the flow on the desired system using the
992
+ requirements specified in `@resources`.
1078
993
 
1079
994
  Parameters
1080
995
  ----------
1081
- packages : Dict[str, str], default: {}
1082
- Packages to use for this step. The key is the name of the package
1083
- and the value is the version to use.
1084
- python : str, optional, default: None
1085
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1086
- that the version used will correspond to the version of the Python interpreter used to start the run.
996
+ cpu : int, default 1
997
+ Number of CPUs required for this step.
998
+ gpu : int, default 0
999
+ Number of GPUs required for this step.
1000
+ disk : int, optional, default None
1001
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
1002
+ memory : int, default 4096
1003
+ Memory size (in MB) required for this step.
1004
+ shared_memory : int, optional, default None
1005
+ The value for the size (in MiB) of the /dev/shm volume for this step.
1006
+ This parameter maps to the `--shm-size` option in Docker.
1087
1007
  """
1088
1008
  ...
1089
1009
 
1090
1010
  @typing.overload
1091
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1011
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1092
1012
  ...
1093
1013
 
1094
1014
  @typing.overload
1095
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1015
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1096
1016
  ...
1097
1017
 
1098
- 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
+ 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):
1099
1019
  """
1100
- Specifies the PyPI packages for the step.
1020
+ Specifies the resources needed when executing this step.
1101
1021
 
1102
- Information in this decorator will augment any
1103
- attributes set in the `@pyi_base` flow-level decorator. Hence,
1104
- you can use `@pypi_base` to set packages required by all
1105
- steps and use `@pypi` to specify step-specific overrides.
1022
+ Use `@resources` to specify the resource requirements
1023
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
1024
+
1025
+ You can choose the compute layer on the command line by executing e.g.
1026
+ ```
1027
+ python myflow.py run --with batch
1028
+ ```
1029
+ or
1030
+ ```
1031
+ python myflow.py run --with kubernetes
1032
+ ```
1033
+ which executes the flow on the desired system using the
1034
+ requirements specified in `@resources`.
1106
1035
 
1107
1036
  Parameters
1108
1037
  ----------
1109
- packages : Dict[str, str], default: {}
1110
- Packages to use for this step. The key is the name of the package
1111
- and the value is the version to use.
1112
- python : str, optional, default: None
1113
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1114
- that the version used will correspond to the version of the Python interpreter used to start the run.
1038
+ cpu : int, default 1
1039
+ Number of CPUs required for this step.
1040
+ gpu : int, default 0
1041
+ Number of GPUs required for this step.
1042
+ disk : int, optional, default None
1043
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
1044
+ memory : int, default 4096
1045
+ Memory size (in MB) required for this step.
1046
+ shared_memory : int, optional, default None
1047
+ The value for the size (in MiB) of the /dev/shm volume for this step.
1048
+ This parameter maps to the `--shm-size` option in Docker.
1049
+ """
1050
+ ...
1051
+
1052
+ 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]]]:
1053
+ """
1054
+ Specifies that this step should execute on Kubernetes.
1055
+
1056
+ Parameters
1057
+ ----------
1058
+ cpu : int, default 1
1059
+ Number of CPUs required for this step. If `@resources` is
1060
+ also present, the maximum value from all decorators is used.
1061
+ memory : int, default 4096
1062
+ Memory size (in MB) required for this step. If
1063
+ `@resources` is also present, the maximum value from all decorators is
1064
+ used.
1065
+ disk : int, default 10240
1066
+ Disk size (in MB) required for this step. If
1067
+ `@resources` is also present, the maximum value from all decorators is
1068
+ used.
1069
+ image : str, optional, default None
1070
+ Docker image to use when launching on Kubernetes. If not specified, and
1071
+ METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
1072
+ not, a default Docker image mapping to the current version of Python is used.
1073
+ image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
1074
+ If given, the imagePullPolicy to be applied to the Docker image of the step.
1075
+ service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
1076
+ Kubernetes service account to use when launching pod in Kubernetes.
1077
+ secrets : List[str], optional, default None
1078
+ Kubernetes secrets to use when launching pod in Kubernetes. These
1079
+ secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
1080
+ in Metaflow configuration.
1081
+ node_selector: Union[Dict[str,str], str], optional, default None
1082
+ Kubernetes node selector(s) to apply to the pod running the task.
1083
+ Can be passed in as a comma separated string of values e.g. "kubernetes.io/os=linux,kubernetes.io/arch=amd64"
1084
+ or as a dictionary {"kubernetes.io/os": "linux", "kubernetes.io/arch": "amd64"}
1085
+ namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
1086
+ Kubernetes namespace to use when launching pod in Kubernetes.
1087
+ gpu : int, optional, default None
1088
+ Number of GPUs required for this step. A value of zero implies that
1089
+ the scheduled node should not have GPUs.
1090
+ gpu_vendor : str, default KUBERNETES_GPU_VENDOR
1091
+ The vendor of the GPUs to be used for this step.
1092
+ tolerations : List[str], default []
1093
+ The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
1094
+ Kubernetes tolerations to use when launching pod in Kubernetes.
1095
+ use_tmpfs : bool, default False
1096
+ This enables an explicit tmpfs mount for this step.
1097
+ tmpfs_tempdir : bool, default True
1098
+ sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
1099
+ tmpfs_size : int, optional, default: None
1100
+ The value for the size (in MiB) of the tmpfs mount for this step.
1101
+ This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
1102
+ memory allocated for this step.
1103
+ tmpfs_path : str, optional, default /metaflow_temp
1104
+ Path to tmpfs mount for this step.
1105
+ persistent_volume_claims : Dict[str, str], optional, default None
1106
+ A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
1107
+ volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
1108
+ shared_memory: int, optional
1109
+ Shared memory size (in MiB) required for this step
1110
+ port: int, optional
1111
+ Port number to specify in the Kubernetes job object
1112
+ compute_pool : str, optional, default None
1113
+ Compute pool to be used for for this step.
1114
+ If not specified, any accessible compute pool within the perimeter is used.
1115
1115
  """
1116
1116
  ...
1117
1117
 
@@ -1154,13 +1154,41 @@ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packag
1154
1154
  """
1155
1155
  ...
1156
1156
 
1157
- def airflow_s3_key_sensor(*, timeout: int, poke_interval: int, mode: str, exponential_backoff: bool, pool: str, soft_fail: bool, name: str, description: str, bucket_key: typing.Union[str, typing.List[str]], bucket_name: str, wildcard_match: bool, aws_conn_id: str, verify: bool) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1157
+ def nim(*, models: "list[NIM]", backend: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1158
1158
  """
1159
- The `@airflow_s3_key_sensor` decorator attaches a Airflow [S3KeySensor](https://airflow.apache.org/docs/apache-airflow-providers-amazon/stable/_api/airflow/providers/amazon/aws/sensors/s3/index.html#airflow.providers.amazon.aws.sensors.s3.S3KeySensor)
1160
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1161
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1162
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
1163
- starts only after all sensors finish.
1159
+ This decorator is used to run NIM containers in Metaflow tasks as sidecars.
1160
+
1161
+ User code call
1162
+ -----------
1163
+ @nim(
1164
+ models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
1165
+ backend='managed'
1166
+ )
1167
+
1168
+ Valid backend options
1169
+ ---------------------
1170
+ - 'managed': Outerbounds selects a compute provider based on the model.
1171
+ - 🚧 'dataplane': Run in your account.
1172
+
1173
+ Valid model options
1174
+ ----------------
1175
+ - 'meta/llama3-8b-instruct': 8B parameter model
1176
+ - 'meta/llama3-70b-instruct': 70B parameter model
1177
+ - Upon request, any model here: https://nvcf.ngc.nvidia.com/functions?filter=nvidia-functions
1178
+
1179
+ Parameters
1180
+ ----------
1181
+ models: list[NIM]
1182
+ List of NIM containers running models in sidecars.
1183
+ backend: str
1184
+ Compute provider to run the NIM container.
1185
+ """
1186
+ ...
1187
+
1188
+ 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]]:
1189
+ """
1190
+ 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.
1191
+ 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.
1164
1192
 
1165
1193
  Parameters
1166
1194
  ----------
@@ -1181,18 +1209,21 @@ def airflow_s3_key_sensor(*, timeout: int, poke_interval: int, mode: str, expone
1181
1209
  Name of the sensor on Airflow
1182
1210
  description : str
1183
1211
  Description of sensor in the Airflow UI
1184
- bucket_key : Union[str, List[str]]
1185
- The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1186
- When it's specified as a full s3:// url, please leave `bucket_name` as None
1187
- bucket_name : str
1188
- Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1189
- When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1190
- wildcard_match : bool
1191
- whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1192
- aws_conn_id : str
1193
- a reference to the s3 connection on Airflow. (Default: None)
1194
- verify : bool
1195
- Whether or not to verify SSL certificates for S3 connection. (Default: None)
1212
+ external_dag_id : str
1213
+ The dag_id that contains the task you want to wait for.
1214
+ external_task_ids : List[str]
1215
+ The list of task_ids that you want to wait for.
1216
+ If None (default value) the sensor waits for the DAG. (Default: None)
1217
+ allowed_states : List[str]
1218
+ Iterable of allowed states, (Default: ['success'])
1219
+ failed_states : List[str]
1220
+ Iterable of failed or dis-allowed states. (Default: None)
1221
+ execution_delta : datetime.timedelta
1222
+ time difference with the previous execution to look at,
1223
+ the default is the same logical date as the current task or DAG. (Default: None)
1224
+ check_existence: bool
1225
+ Set to True to check if the external task exists or check if
1226
+ the DAG to wait for exists. (Default: True)
1196
1227
  """
1197
1228
  ...
1198
1229
 
@@ -1340,6 +1371,24 @@ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: t
1340
1371
  """
1341
1372
  ...
1342
1373
 
1374
+ def project(*, name: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1375
+ """
1376
+ Specifies what flows belong to the same project.
1377
+
1378
+ A project-specific namespace is created for all flows that
1379
+ use the same `@project(name)`.
1380
+
1381
+ Parameters
1382
+ ----------
1383
+ name : str
1384
+ Project name. Make sure that the name is unique amongst all
1385
+ projects that use the same production scheduler. The name may
1386
+ contain only lowercase alphanumeric characters and underscores.
1387
+
1388
+
1389
+ """
1390
+ ...
1391
+
1343
1392
  @typing.overload
1344
1393
  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]]:
1345
1394
  """
@@ -1389,6 +1438,48 @@ def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly:
1389
1438
  """
1390
1439
  ...
1391
1440
 
1441
+ def airflow_s3_key_sensor(*, timeout: int, poke_interval: int, mode: str, exponential_backoff: bool, pool: str, soft_fail: bool, name: str, description: str, bucket_key: typing.Union[str, typing.List[str]], bucket_name: str, wildcard_match: bool, aws_conn_id: str, verify: bool) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1442
+ """
1443
+ The `@airflow_s3_key_sensor` decorator attaches a Airflow [S3KeySensor](https://airflow.apache.org/docs/apache-airflow-providers-amazon/stable/_api/airflow/providers/amazon/aws/sensors/s3/index.html#airflow.providers.amazon.aws.sensors.s3.S3KeySensor)
1444
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1445
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1446
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
1447
+ starts only after all sensors finish.
1448
+
1449
+ Parameters
1450
+ ----------
1451
+ timeout : int
1452
+ Time, in seconds before the task times out and fails. (Default: 3600)
1453
+ poke_interval : int
1454
+ Time in seconds that the job should wait in between each try. (Default: 60)
1455
+ mode : str
1456
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1457
+ exponential_backoff : bool
1458
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1459
+ pool : str
1460
+ the slot pool this task should run in,
1461
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1462
+ soft_fail : bool
1463
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1464
+ name : str
1465
+ Name of the sensor on Airflow
1466
+ description : str
1467
+ Description of sensor in the Airflow UI
1468
+ bucket_key : Union[str, List[str]]
1469
+ The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1470
+ When it's specified as a full s3:// url, please leave `bucket_name` as None
1471
+ bucket_name : str
1472
+ Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1473
+ When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1474
+ wildcard_match : bool
1475
+ whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1476
+ aws_conn_id : str
1477
+ a reference to the s3 connection on Airflow. (Default: None)
1478
+ verify : bool
1479
+ Whether or not to verify SSL certificates for S3 connection. (Default: None)
1480
+ """
1481
+ ...
1482
+
1392
1483
  @typing.overload
1393
1484
  def trigger_on_finish(*, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1394
1485
  """
@@ -1492,97 +1583,6 @@ def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *
1492
1583
  """
1493
1584
  ...
1494
1585
 
1495
- def project(*, name: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1496
- """
1497
- Specifies what flows belong to the same project.
1498
-
1499
- A project-specific namespace is created for all flows that
1500
- use the same `@project(name)`.
1501
-
1502
- Parameters
1503
- ----------
1504
- name : str
1505
- Project name. Make sure that the name is unique amongst all
1506
- projects that use the same production scheduler. The name may
1507
- contain only lowercase alphanumeric characters and underscores.
1508
-
1509
-
1510
- """
1511
- ...
1512
-
1513
- 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]]:
1514
- """
1515
- 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.
1516
- 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.
1517
-
1518
- Parameters
1519
- ----------
1520
- timeout : int
1521
- Time, in seconds before the task times out and fails. (Default: 3600)
1522
- poke_interval : int
1523
- Time in seconds that the job should wait in between each try. (Default: 60)
1524
- mode : str
1525
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1526
- exponential_backoff : bool
1527
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1528
- pool : str
1529
- the slot pool this task should run in,
1530
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1531
- soft_fail : bool
1532
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1533
- name : str
1534
- Name of the sensor on Airflow
1535
- description : str
1536
- Description of sensor in the Airflow UI
1537
- external_dag_id : str
1538
- The dag_id that contains the task you want to wait for.
1539
- external_task_ids : List[str]
1540
- The list of task_ids that you want to wait for.
1541
- If None (default value) the sensor waits for the DAG. (Default: None)
1542
- allowed_states : List[str]
1543
- Iterable of allowed states, (Default: ['success'])
1544
- failed_states : List[str]
1545
- Iterable of failed or dis-allowed states. (Default: None)
1546
- execution_delta : datetime.timedelta
1547
- time difference with the previous execution to look at,
1548
- the default is the same logical date as the current task or DAG. (Default: None)
1549
- check_existence: bool
1550
- Set to True to check if the external task exists or check if
1551
- the DAG to wait for exists. (Default: True)
1552
- """
1553
- ...
1554
-
1555
- def nim(*, models: "list[NIM]", backend: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1556
- """
1557
- This decorator is used to run NIM containers in Metaflow tasks as sidecars.
1558
-
1559
- User code call
1560
- -----------
1561
- @nim(
1562
- models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
1563
- backend='managed'
1564
- )
1565
-
1566
- Valid backend options
1567
- ---------------------
1568
- - 'managed': Outerbounds selects a compute provider based on the model.
1569
- - 🚧 'dataplane': Run in your account.
1570
-
1571
- Valid model options
1572
- ----------------
1573
- - 'meta/llama3-8b-instruct': 8B parameter model
1574
- - 'meta/llama3-70b-instruct': 70B parameter model
1575
- - Upon request, any model here: https://nvcf.ngc.nvidia.com/functions?filter=nvidia-functions
1576
-
1577
- Parameters
1578
- ----------
1579
- models: list[NIM]
1580
- List of NIM containers running models in sidecars.
1581
- backend: str
1582
- Compute provider to run the NIM container.
1583
- """
1584
- ...
1585
-
1586
1586
  def namespace(ns: typing.Optional[str]) -> typing.Optional[str]:
1587
1587
  """
1588
1588
  Switch namespace to the one provided.