ob-metaflow-stubs 6.0.3.158__py2.py3-none-any.whl → 6.0.3.159__py2.py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (200) hide show
  1. metaflow-stubs/__init__.pyi +768 -768
  2. metaflow-stubs/cards.pyi +1 -1
  3. metaflow-stubs/cli.pyi +1 -1
  4. metaflow-stubs/cli_components/__init__.pyi +1 -1
  5. metaflow-stubs/cli_components/utils.pyi +1 -1
  6. metaflow-stubs/client/__init__.pyi +1 -1
  7. metaflow-stubs/client/core.pyi +4 -4
  8. metaflow-stubs/client/filecache.pyi +1 -1
  9. metaflow-stubs/events.pyi +1 -1
  10. metaflow-stubs/exception.pyi +1 -1
  11. metaflow-stubs/flowspec.pyi +3 -3
  12. metaflow-stubs/generated_for.txt +1 -1
  13. metaflow-stubs/includefile.pyi +2 -2
  14. metaflow-stubs/info_file.pyi +1 -1
  15. metaflow-stubs/metadata_provider/__init__.pyi +1 -1
  16. metaflow-stubs/metadata_provider/heartbeat.pyi +1 -1
  17. metaflow-stubs/metadata_provider/metadata.pyi +1 -1
  18. metaflow-stubs/metadata_provider/util.pyi +1 -1
  19. metaflow-stubs/metaflow_config.pyi +1 -1
  20. metaflow-stubs/metaflow_current.pyi +119 -119
  21. metaflow-stubs/mf_extensions/__init__.pyi +1 -1
  22. metaflow-stubs/mf_extensions/obcheckpoint/__init__.pyi +1 -1
  23. metaflow-stubs/mf_extensions/obcheckpoint/plugins/__init__.pyi +1 -1
  24. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/__init__.pyi +1 -1
  25. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/__init__.pyi +1 -1
  26. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/async_cards.pyi +1 -1
  27. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/deco_injection_mixin.pyi +1 -1
  28. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/extra_components.pyi +2 -2
  29. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/__init__.pyi +1 -1
  30. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/__init__.pyi +1 -1
  31. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/checkpoint_lister.pyi +3 -3
  32. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/lineage_card.pyi +1 -1
  33. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/checkpoint_storage.pyi +3 -3
  34. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/constructors.pyi +1 -1
  35. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/core.pyi +3 -3
  36. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/decorator.pyi +3 -3
  37. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/exceptions.pyi +1 -1
  38. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/final_api.pyi +2 -2
  39. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/lineage.pyi +1 -1
  40. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/__init__.pyi +1 -1
  41. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/context.pyi +1 -1
  42. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/core.pyi +2 -2
  43. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/decorator.pyi +1 -1
  44. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/exceptions.pyi +1 -1
  45. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/task_utils.pyi +1 -1
  46. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/utils.pyi +1 -1
  47. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastructures.pyi +1 -1
  48. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/exceptions.pyi +1 -1
  49. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/__init__.pyi +1 -1
  50. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/decorator.pyi +1 -1
  51. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/__init__.pyi +1 -1
  52. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/core.pyi +2 -2
  53. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/exceptions.pyi +1 -1
  54. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/model_storage.pyi +2 -2
  55. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/__init__.pyi +1 -1
  56. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/flowspec_utils.pyi +1 -1
  57. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/general.pyi +1 -1
  58. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/identity_utils.pyi +1 -1
  59. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/__init__.pyi +1 -1
  60. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/base.pyi +1 -1
  61. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/tar.pyi +1 -1
  62. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/tar_utils.pyi +1 -1
  63. metaflow-stubs/mf_extensions/outerbounds/__init__.pyi +1 -1
  64. metaflow-stubs/mf_extensions/outerbounds/plugins/__init__.pyi +1 -1
  65. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/__init__.pyi +1 -1
  66. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/injector.pyi +1 -1
  67. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/__init__.pyi +1 -1
  68. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/ollama.pyi +1 -1
  69. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/__init__.pyi +1 -1
  70. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/snowflake.pyi +1 -1
  71. metaflow-stubs/mf_extensions/outerbounds/profilers/__init__.pyi +1 -1
  72. metaflow-stubs/mf_extensions/outerbounds/profilers/gpu.pyi +1 -1
  73. metaflow-stubs/mf_extensions/outerbounds/remote_config.pyi +1 -1
  74. metaflow-stubs/mf_extensions/outerbounds/toplevel/__init__.pyi +1 -1
  75. metaflow-stubs/mf_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.pyi +1 -1
  76. metaflow-stubs/multicore_utils.pyi +1 -1
  77. metaflow-stubs/parameters.pyi +2 -2
  78. metaflow-stubs/plugins/__init__.pyi +10 -10
  79. metaflow-stubs/plugins/airflow/__init__.pyi +1 -1
  80. metaflow-stubs/plugins/airflow/airflow_utils.pyi +1 -1
  81. metaflow-stubs/plugins/airflow/exception.pyi +1 -1
  82. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +1 -1
  83. metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi +1 -1
  84. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +1 -1
  85. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +1 -1
  86. metaflow-stubs/plugins/argo/__init__.pyi +1 -1
  87. metaflow-stubs/plugins/argo/argo_client.pyi +1 -1
  88. metaflow-stubs/plugins/argo/argo_events.pyi +1 -1
  89. metaflow-stubs/plugins/argo/argo_workflows.pyi +2 -2
  90. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +2 -2
  91. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +1 -1
  92. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +2 -2
  93. metaflow-stubs/plugins/aws/__init__.pyi +1 -1
  94. metaflow-stubs/plugins/aws/aws_client.pyi +1 -1
  95. metaflow-stubs/plugins/aws/aws_utils.pyi +1 -1
  96. metaflow-stubs/plugins/aws/batch/__init__.pyi +1 -1
  97. metaflow-stubs/plugins/aws/batch/batch.pyi +1 -1
  98. metaflow-stubs/plugins/aws/batch/batch_client.pyi +1 -1
  99. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +1 -1
  100. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +1 -1
  101. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +2 -2
  102. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +1 -1
  103. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +1 -1
  104. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +1 -1
  105. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +1 -1
  106. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +1 -1
  107. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer.pyi +2 -2
  108. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer_objects.pyi +2 -2
  109. metaflow-stubs/plugins/azure/__init__.pyi +1 -1
  110. metaflow-stubs/plugins/azure/azure_credential.pyi +1 -1
  111. metaflow-stubs/plugins/azure/azure_exceptions.pyi +1 -1
  112. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +2 -2
  113. metaflow-stubs/plugins/azure/azure_utils.pyi +1 -1
  114. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +1 -1
  115. metaflow-stubs/plugins/azure/includefile_support.pyi +1 -1
  116. metaflow-stubs/plugins/cards/__init__.pyi +5 -5
  117. metaflow-stubs/plugins/cards/card_client.pyi +1 -1
  118. metaflow-stubs/plugins/cards/card_creator.pyi +1 -1
  119. metaflow-stubs/plugins/cards/card_datastore.pyi +1 -1
  120. metaflow-stubs/plugins/cards/card_decorator.pyi +1 -1
  121. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +1 -1
  122. metaflow-stubs/plugins/cards/card_modules/basic.pyi +2 -2
  123. metaflow-stubs/plugins/cards/card_modules/card.pyi +1 -1
  124. metaflow-stubs/plugins/cards/card_modules/components.pyi +2 -2
  125. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +1 -1
  126. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +1 -1
  127. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +1 -1
  128. metaflow-stubs/plugins/cards/card_resolver.pyi +1 -1
  129. metaflow-stubs/plugins/cards/component_serializer.pyi +1 -1
  130. metaflow-stubs/plugins/cards/exception.pyi +1 -1
  131. metaflow-stubs/plugins/catch_decorator.pyi +1 -1
  132. metaflow-stubs/plugins/datatools/__init__.pyi +1 -1
  133. metaflow-stubs/plugins/datatools/local.pyi +1 -1
  134. metaflow-stubs/plugins/datatools/s3/__init__.pyi +1 -1
  135. metaflow-stubs/plugins/datatools/s3/s3.pyi +2 -2
  136. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +1 -1
  137. metaflow-stubs/plugins/datatools/s3/s3util.pyi +1 -1
  138. metaflow-stubs/plugins/debug_logger.pyi +1 -1
  139. metaflow-stubs/plugins/debug_monitor.pyi +1 -1
  140. metaflow-stubs/plugins/environment_decorator.pyi +1 -1
  141. metaflow-stubs/plugins/events_decorator.pyi +1 -1
  142. metaflow-stubs/plugins/frameworks/__init__.pyi +1 -1
  143. metaflow-stubs/plugins/frameworks/pytorch.pyi +1 -1
  144. metaflow-stubs/plugins/gcp/__init__.pyi +1 -1
  145. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +2 -2
  146. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +1 -1
  147. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +1 -1
  148. metaflow-stubs/plugins/gcp/gs_utils.pyi +1 -1
  149. metaflow-stubs/plugins/gcp/includefile_support.pyi +1 -1
  150. metaflow-stubs/plugins/kubernetes/__init__.pyi +1 -1
  151. metaflow-stubs/plugins/kubernetes/kube_utils.pyi +1 -1
  152. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +1 -1
  153. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +1 -1
  154. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +1 -1
  155. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +1 -1
  156. metaflow-stubs/plugins/kubernetes/spot_monitor_sidecar.pyi +1 -1
  157. metaflow-stubs/plugins/ollama/__init__.pyi +1 -1
  158. metaflow-stubs/plugins/parallel_decorator.pyi +1 -1
  159. metaflow-stubs/plugins/perimeters.pyi +1 -1
  160. metaflow-stubs/plugins/project_decorator.pyi +1 -1
  161. metaflow-stubs/plugins/pypi/__init__.pyi +2 -2
  162. metaflow-stubs/plugins/pypi/conda_decorator.pyi +1 -1
  163. metaflow-stubs/plugins/pypi/conda_environment.pyi +2 -2
  164. metaflow-stubs/plugins/pypi/parsers.pyi +1 -1
  165. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +1 -1
  166. metaflow-stubs/plugins/pypi/pypi_environment.pyi +1 -1
  167. metaflow-stubs/plugins/pypi/utils.pyi +1 -1
  168. metaflow-stubs/plugins/resources_decorator.pyi +1 -1
  169. metaflow-stubs/plugins/retry_decorator.pyi +1 -1
  170. metaflow-stubs/plugins/secrets/__init__.pyi +2 -2
  171. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +2 -2
  172. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +1 -1
  173. metaflow-stubs/plugins/snowflake/__init__.pyi +1 -1
  174. metaflow-stubs/plugins/storage_executor.pyi +1 -1
  175. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +1 -1
  176. metaflow-stubs/plugins/timeout_decorator.pyi +1 -1
  177. metaflow-stubs/profilers/__init__.pyi +1 -1
  178. metaflow-stubs/pylint_wrapper.pyi +1 -1
  179. metaflow-stubs/runner/__init__.pyi +1 -1
  180. metaflow-stubs/runner/deployer.pyi +29 -29
  181. metaflow-stubs/runner/deployer_impl.pyi +1 -1
  182. metaflow-stubs/runner/metaflow_runner.pyi +2 -2
  183. metaflow-stubs/runner/nbdeploy.pyi +1 -1
  184. metaflow-stubs/runner/nbrun.pyi +1 -1
  185. metaflow-stubs/runner/subprocess_manager.pyi +1 -1
  186. metaflow-stubs/runner/utils.pyi +2 -2
  187. metaflow-stubs/system/__init__.pyi +1 -1
  188. metaflow-stubs/system/system_logger.pyi +1 -1
  189. metaflow-stubs/system/system_monitor.pyi +1 -1
  190. metaflow-stubs/tagging_util.pyi +1 -1
  191. metaflow-stubs/tuple_util.pyi +1 -1
  192. metaflow-stubs/user_configs/__init__.pyi +1 -1
  193. metaflow-stubs/user_configs/config_decorators.pyi +5 -5
  194. metaflow-stubs/user_configs/config_options.pyi +2 -2
  195. metaflow-stubs/user_configs/config_parameters.pyi +6 -6
  196. {ob_metaflow_stubs-6.0.3.158.dist-info → ob_metaflow_stubs-6.0.3.159.dist-info}/METADATA +1 -1
  197. ob_metaflow_stubs-6.0.3.159.dist-info/RECORD +200 -0
  198. ob_metaflow_stubs-6.0.3.158.dist-info/RECORD +0 -200
  199. {ob_metaflow_stubs-6.0.3.158.dist-info → ob_metaflow_stubs-6.0.3.159.dist-info}/WHEEL +0 -0
  200. {ob_metaflow_stubs-6.0.3.158.dist-info → ob_metaflow_stubs-6.0.3.159.dist-info}/top_level.txt +0 -0
@@ -1,15 +1,15 @@
1
1
  ######################################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
3
  # MF version: 2.15.7.2+obcheckpoint(0.2.1);ob(v1) #
4
- # Generated on 2025-04-17T20:45:30.855915 #
4
+ # Generated on 2025-04-18T04:20:36.937065 #
5
5
  ######################################################################################################
6
6
 
7
7
  from __future__ import annotations
8
8
 
9
9
  import typing
10
10
  if typing.TYPE_CHECKING:
11
- import typing
12
11
  import datetime
12
+ import typing
13
13
  FlowSpecDerived = typing.TypeVar("FlowSpecDerived", bound="FlowSpec", contravariant=False, covariant=False)
14
14
  StepFlag = typing.NewType("StepFlag", bool)
15
15
 
@@ -36,16 +36,16 @@ from .user_configs.config_parameters import config_expr as config_expr
36
36
  from .user_configs.config_decorators import CustomFlowDecorator as CustomFlowDecorator
37
37
  from .user_configs.config_decorators import CustomStepDecorator as CustomStepDecorator
38
38
  from . import cards as cards
39
- from . import tuple_util as tuple_util
40
39
  from . import events as events
40
+ from . import tuple_util as tuple_util
41
41
  from . import runner as runner
42
42
  from . import plugins as plugins
43
43
  from .mf_extensions.outerbounds.toplevel.global_aliases_for_metaflow_package import S3 as S3
44
44
  from . import includefile as includefile
45
45
  from .includefile import IncludeFile as IncludeFile
46
- from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
47
- from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
48
46
  from .plugins.pypi.parsers import conda_environment_yml_parser as conda_environment_yml_parser
47
+ from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
48
+ from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
49
49
  from . import client as client
50
50
  from .client.core import namespace as namespace
51
51
  from .client.core import get_namespace as get_namespace
@@ -150,136 +150,145 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
150
150
  """
151
151
  ...
152
152
 
153
- @typing.overload
154
- 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]]]:
153
+ def app_deploy(*, app_port: int, app_name: 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]]]:
155
154
  """
156
- Specifies that the step will success under all circumstances.
157
-
158
- The decorator will create an optional artifact, specified by `var`, which
159
- contains the exception raised. You can use it to detect the presence
160
- of errors, indicating that all happy-path artifacts produced by the step
161
- are missing.
155
+ Specifies that this step is used to deploy an instance of the app.
156
+ Requires that self.app_name, self.app_port, self.entrypoint and self.deployDir is set.
162
157
 
163
158
 
164
159
  Parameters
165
160
  ----------
166
- var : str, optional, default None
167
- Name of the artifact in which to store the caught exception.
168
- If not specified, the exception is not stored.
169
- print_exception : bool, default True
170
- Determines whether or not the exception is printed to
171
- stdout when caught.
161
+ app_port : int
162
+ Number of GPUs to use.
163
+ app_name : str
164
+ Name of the app to deploy.
172
165
  """
173
166
  ...
174
167
 
175
- @typing.overload
176
- def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
177
- ...
178
-
179
- @typing.overload
180
- def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
181
- ...
182
-
183
- 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):
168
+ def ollama(*, models: "list[Ollama]", backend: 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]]]:
184
169
  """
185
- Specifies that the step will success under all circumstances.
170
+ This decorator is used to run Ollama APIs as Metaflow task sidecars.
186
171
 
187
- The decorator will create an optional artifact, specified by `var`, which
188
- contains the exception raised. You can use it to detect the presence
189
- of errors, indicating that all happy-path artifacts produced by the step
190
- are missing.
172
+ User code call
173
+ -----------
174
+ @ollama(
175
+ models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
176
+ backend='local'
177
+ )
178
+
179
+ Valid backend options
180
+ ---------------------
181
+ - 'local': Run as a separate process on the local task machine.
182
+ - (TODO) 'managed': Outerbounds hosts and selects compute provider.
183
+ - (TODO) 'remote': Spin up separate instance to serve Ollama models.
184
+
185
+ Valid model options
186
+ ----------------
187
+ - 'llama3.2'
188
+ - 'llama3.3'
189
+ - any model here https://ollama.com/search
191
190
 
192
191
 
193
192
  Parameters
194
193
  ----------
195
- var : str, optional, default None
196
- Name of the artifact in which to store the caught exception.
197
- If not specified, the exception is not stored.
198
- print_exception : bool, default True
199
- Determines whether or not the exception is printed to
200
- stdout when caught.
194
+ models: list[Ollama]
195
+ List of Ollama containers running models in sidecars.
196
+ backend: str
197
+ Determines where and how to run the Ollama process.
201
198
  """
202
199
  ...
203
200
 
204
201
  @typing.overload
205
- 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]]]:
202
+ 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]]]:
206
203
  """
207
- Creates a human-readable report, a Metaflow Card, after this step completes.
208
-
209
- Note that you may add multiple `@card` decorators in a step with different parameters.
204
+ Specifies secrets to be retrieved and injected as environment variables prior to
205
+ the execution of a step.
210
206
 
211
207
 
212
208
  Parameters
213
209
  ----------
214
- type : str, default 'default'
215
- Card type.
216
- id : str, optional, default None
217
- If multiple cards are present, use this id to identify this card.
218
- options : Dict[str, Any], default {}
219
- Options passed to the card. The contents depend on the card type.
220
- timeout : int, default 45
221
- Interrupt reporting if it takes more than this many seconds.
210
+ sources : List[Union[str, Dict[str, Any]]], default: []
211
+ List of secret specs, defining how the secrets are to be retrieved
222
212
  """
223
213
  ...
224
214
 
225
215
  @typing.overload
226
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
216
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
227
217
  ...
228
218
 
229
219
  @typing.overload
230
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
220
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
231
221
  ...
232
222
 
233
- 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):
223
+ 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]]] = []):
234
224
  """
235
- Creates a human-readable report, a Metaflow Card, after this step completes.
236
-
237
- Note that you may add multiple `@card` decorators in a step with different parameters.
225
+ Specifies secrets to be retrieved and injected as environment variables prior to
226
+ the execution of a step.
238
227
 
239
228
 
240
229
  Parameters
241
230
  ----------
242
- type : str, default 'default'
243
- Card type.
244
- id : str, optional, default None
245
- If multiple cards are present, use this id to identify this card.
246
- options : Dict[str, Any], default {}
247
- Options passed to the card. The contents depend on the card type.
248
- timeout : int, default 45
249
- Interrupt reporting if it takes more than this many seconds.
231
+ sources : List[Union[str, Dict[str, Any]]], default: []
232
+ List of secret specs, defining how the secrets are to be retrieved
250
233
  """
251
234
  ...
252
235
 
253
- def ollama(*, models: "list[Ollama]", backend: 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]]]:
236
+ @typing.overload
237
+ 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]]]:
254
238
  """
255
- This decorator is used to run Ollama APIs as Metaflow task sidecars.
239
+ Specifies the Conda environment for the step.
256
240
 
257
- User code call
258
- -----------
259
- @ollama(
260
- models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
261
- backend='local'
262
- )
241
+ Information in this decorator will augment any
242
+ attributes set in the `@conda_base` flow-level decorator. Hence,
243
+ you can use `@conda_base` to set packages required by all
244
+ steps and use `@conda` to specify step-specific overrides.
263
245
 
264
- Valid backend options
265
- ---------------------
266
- - 'local': Run as a separate process on the local task machine.
267
- - (TODO) 'managed': Outerbounds hosts and selects compute provider.
268
- - (TODO) 'remote': Spin up separate instance to serve Ollama models.
269
246
 
270
- Valid model options
271
- ----------------
272
- - 'llama3.2'
273
- - 'llama3.3'
274
- - any model here https://ollama.com/search
247
+ Parameters
248
+ ----------
249
+ packages : Dict[str, str], default {}
250
+ Packages to use for this step. The key is the name of the package
251
+ and the value is the version to use.
252
+ libraries : Dict[str, str], default {}
253
+ Supported for backward compatibility. When used with packages, packages will take precedence.
254
+ python : str, optional, default None
255
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
256
+ that the version used will correspond to the version of the Python interpreter used to start the run.
257
+ disabled : bool, default False
258
+ If set to True, disables @conda.
259
+ """
260
+ ...
261
+
262
+ @typing.overload
263
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
264
+ ...
265
+
266
+ @typing.overload
267
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
268
+ ...
269
+
270
+ 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):
271
+ """
272
+ Specifies the Conda environment for the step.
273
+
274
+ Information in this decorator will augment any
275
+ attributes set in the `@conda_base` flow-level decorator. Hence,
276
+ you can use `@conda_base` to set packages required by all
277
+ steps and use `@conda` to specify step-specific overrides.
275
278
 
276
279
 
277
280
  Parameters
278
281
  ----------
279
- models: list[Ollama]
280
- List of Ollama containers running models in sidecars.
281
- backend: str
282
- Determines where and how to run the Ollama process.
282
+ packages : Dict[str, str], default {}
283
+ Packages to use for this step. The key is the name of the package
284
+ and the value is the version to use.
285
+ libraries : Dict[str, str], default {}
286
+ Supported for backward compatibility. When used with packages, packages will take precedence.
287
+ python : str, optional, default None
288
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
289
+ that the version used will correspond to the version of the Python interpreter used to start the run.
290
+ disabled : bool, default False
291
+ If set to True, disables @conda.
283
292
  """
284
293
  ...
285
294
 
@@ -360,260 +369,170 @@ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: ty
360
369
  """
361
370
  ...
362
371
 
363
- def nim(*, models: "list[NIM]", backend: str, queue_timeout: int) -> 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]]]:
372
+ @typing.overload
373
+ def resources(*, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
364
374
  """
365
- This decorator is used to run NIM containers in Metaflow tasks as sidecars.
366
-
367
- User code call
368
- -----------
369
- @nim(
370
- models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
371
- backend='managed'
372
- )
375
+ Specifies the resources needed when executing this step.
373
376
 
374
- Valid backend options
375
- ---------------------
376
- - 'managed': Outerbounds selects a compute provider based on the model.
377
+ Use `@resources` to specify the resource requirements
378
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
377
379
 
378
- Valid model options
379
- ----------------
380
- - 'meta/llama3-8b-instruct': 8B parameter model
381
- - 'meta/llama3-70b-instruct': 70B parameter model
382
- - any model here: https://nvcf.ngc.nvidia.com/functions?filter=nvidia-functions
380
+ You can choose the compute layer on the command line by executing e.g.
381
+ ```
382
+ python myflow.py run --with batch
383
+ ```
384
+ or
385
+ ```
386
+ python myflow.py run --with kubernetes
387
+ ```
388
+ which executes the flow on the desired system using the
389
+ requirements specified in `@resources`.
383
390
 
384
391
 
385
392
  Parameters
386
393
  ----------
387
- models: list[NIM]
388
- List of NIM containers running models in sidecars.
389
- backend: str
390
- Compute provider to run the NIM container.
391
- queue_timeout : int
392
- Time to keep the job in NVCF's queue.
394
+ cpu : int, default 1
395
+ Number of CPUs required for this step.
396
+ gpu : int, optional, default None
397
+ Number of GPUs required for this step.
398
+ disk : int, optional, default None
399
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
400
+ memory : int, default 4096
401
+ Memory size (in MB) required for this step.
402
+ shared_memory : int, optional, default None
403
+ The value for the size (in MiB) of the /dev/shm volume for this step.
404
+ This parameter maps to the `--shm-size` option in Docker.
393
405
  """
394
406
  ...
395
407
 
396
408
  @typing.overload
397
- 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]]]:
409
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
410
+ ...
411
+
412
+ @typing.overload
413
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
414
+ ...
415
+
416
+ def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None):
398
417
  """
399
- Enables checkpointing for a step.
418
+ Specifies the resources needed when executing this step.
419
+
420
+ Use `@resources` to specify the resource requirements
421
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
400
422
 
423
+ You can choose the compute layer on the command line by executing e.g.
424
+ ```
425
+ python myflow.py run --with batch
426
+ ```
427
+ or
428
+ ```
429
+ python myflow.py run --with kubernetes
430
+ ```
431
+ which executes the flow on the desired system using the
432
+ requirements specified in `@resources`.
401
433
 
402
434
 
403
435
  Parameters
404
436
  ----------
405
- load_policy : str, default: "fresh"
406
- The policy for loading the checkpoint. The following policies are supported:
407
- - "eager": Loads the the latest available checkpoint within the namespace.
408
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
409
- will be loaded at the start of the task.
410
- - "none": Do not load any checkpoint
411
- - "fresh": Loads the lastest checkpoint created within the running Task.
412
- This mode helps loading checkpoints across various retry attempts of the same task.
413
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
414
- created within the task will be loaded when the task is retries execution on failure.
415
-
416
- temp_dir_root : str, default: None
417
- The root directory under which `current.checkpoint.directory` will be created.
437
+ cpu : int, default 1
438
+ Number of CPUs required for this step.
439
+ gpu : int, optional, default None
440
+ Number of GPUs required for this step.
441
+ disk : int, optional, default None
442
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
443
+ memory : int, default 4096
444
+ Memory size (in MB) required for this step.
445
+ shared_memory : int, optional, default None
446
+ The value for the size (in MiB) of the /dev/shm volume for this step.
447
+ This parameter maps to the `--shm-size` option in Docker.
418
448
  """
419
449
  ...
420
450
 
421
451
  @typing.overload
422
- def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
423
- ...
424
-
425
- @typing.overload
426
- def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
427
- ...
428
-
429
- 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):
452
+ 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]]]:
430
453
  """
431
- Enables checkpointing for a step.
432
-
454
+ Specifies environment variables to be set prior to the execution of a step.
433
455
 
434
456
 
435
457
  Parameters
436
458
  ----------
437
- load_policy : str, default: "fresh"
438
- The policy for loading the checkpoint. The following policies are supported:
439
- - "eager": Loads the the latest available checkpoint within the namespace.
440
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
441
- will be loaded at the start of the task.
442
- - "none": Do not load any checkpoint
443
- - "fresh": Loads the lastest checkpoint created within the running Task.
444
- This mode helps loading checkpoints across various retry attempts of the same task.
445
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
446
- created within the task will be loaded when the task is retries execution on failure.
447
-
448
- temp_dir_root : str, default: None
449
- The root directory under which `current.checkpoint.directory` will be created.
459
+ vars : Dict[str, str], default {}
460
+ Dictionary of environment variables to set.
450
461
  """
451
462
  ...
452
463
 
453
464
  @typing.overload
454
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
455
- """
456
- Internal decorator to support Fast bakery
457
- """
465
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
458
466
  ...
459
467
 
460
468
  @typing.overload
461
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
462
- ...
463
-
464
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
465
- """
466
- Internal decorator to support Fast bakery
467
- """
469
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
468
470
  ...
469
471
 
470
- def huggingface_hub(*, temp_dir_root: typing.Optional[str] = None, load: typing.Union[typing.List[str], typing.List[typing.Tuple[typing.Dict, str]], typing.List[typing.Tuple[str, str]], typing.List[typing.Dict], 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]]]:
472
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
471
473
  """
472
- Decorator that helps cache, version and store models/datasets from huggingface hub.
474
+ Specifies environment variables to be set prior to the execution of a step.
473
475
 
474
476
 
475
477
  Parameters
476
478
  ----------
477
- temp_dir_root : str, optional
478
- The root directory that will hold the temporary directory where objects will be downloaded.
479
-
480
- load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
481
- The list of repos (models/datasets) to load.
482
-
483
- Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
484
-
485
- - If repo (model/dataset) is not found in the datastore:
486
- - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
487
- - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
488
- - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
489
-
490
- - If repo is found in the datastore:
491
- - Loads it directly from datastore to local path (can be temporary directory or specified path)
479
+ vars : Dict[str, str], default {}
480
+ Dictionary of environment variables to set.
492
481
  """
493
482
  ...
494
483
 
495
484
  @typing.overload
496
- 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]]]:
485
+ 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]]]:
497
486
  """
498
- Enables loading / saving of models within a step.
499
-
500
-
501
-
502
- Parameters
503
- ----------
504
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
505
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
506
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
507
- - `current.checkpoint`
508
- - `current.model`
509
- - `current.huggingface_hub`
510
-
511
- 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
512
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
513
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
487
+ Specifies the number of times the task corresponding
488
+ to a step needs to be retried.
514
489
 
515
- temp_dir_root : str, default: None
516
- The root directory under which `current.model.loaded` will store loaded models
517
- """
518
- ...
519
-
520
- @typing.overload
521
- def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
522
- ...
523
-
524
- @typing.overload
525
- def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
526
- ...
527
-
528
- 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):
529
- """
530
- Enables loading / saving of models within a step.
490
+ This decorator is useful for handling transient errors, such as networking issues.
491
+ If your task contains operations that can't be retried safely, e.g. database updates,
492
+ it is advisable to annotate it with `@retry(times=0)`.
531
493
 
494
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
495
+ decorator will execute a no-op task after all retries have been exhausted,
496
+ ensuring that the flow execution can continue.
532
497
 
533
498
 
534
499
  Parameters
535
500
  ----------
536
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
537
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
538
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
539
- - `current.checkpoint`
540
- - `current.model`
541
- - `current.huggingface_hub`
542
-
543
- 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
544
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
545
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
546
-
547
- temp_dir_root : str, default: None
548
- The root directory under which `current.model.loaded` will store loaded models
501
+ times : int, default 3
502
+ Number of times to retry this task.
503
+ minutes_between_retries : int, default 2
504
+ Number of minutes between retries.
549
505
  """
550
506
  ...
551
507
 
552
508
  @typing.overload
553
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
554
- """
555
- Decorator prototype for all step decorators. This function gets specialized
556
- and imported for all decorators types by _import_plugin_decorators().
557
- """
509
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
558
510
  ...
559
511
 
560
512
  @typing.overload
561
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
562
- ...
563
-
564
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
565
- """
566
- Decorator prototype for all step decorators. This function gets specialized
567
- and imported for all decorators types by _import_plugin_decorators().
568
- """
569
- ...
570
-
571
- def nvidia(*, gpu: int, gpu_type: str, queue_timeout: int) -> 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]]]:
572
- """
573
- Specifies that this step should execute on DGX cloud.
574
-
575
-
576
- Parameters
577
- ----------
578
- gpu : int
579
- Number of GPUs to use.
580
- gpu_type : str
581
- Type of Nvidia GPU to use.
582
- queue_timeout : int
583
- Time to keep the job in NVCF's queue.
584
- """
513
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
585
514
  ...
586
515
 
587
- @typing.overload
588
- 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]]]:
516
+ 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):
589
517
  """
590
- Specifies environment variables to be set prior to the execution of a step.
518
+ Specifies the number of times the task corresponding
519
+ to a step needs to be retried.
591
520
 
521
+ This decorator is useful for handling transient errors, such as networking issues.
522
+ If your task contains operations that can't be retried safely, e.g. database updates,
523
+ it is advisable to annotate it with `@retry(times=0)`.
592
524
 
593
- Parameters
594
- ----------
595
- vars : Dict[str, str], default {}
596
- Dictionary of environment variables to set.
597
- """
598
- ...
599
-
600
- @typing.overload
601
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
602
- ...
603
-
604
- @typing.overload
605
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
606
- ...
607
-
608
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
609
- """
610
- Specifies environment variables to be set prior to the execution of a step.
525
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
526
+ decorator will execute a no-op task after all retries have been exhausted,
527
+ ensuring that the flow execution can continue.
611
528
 
612
529
 
613
530
  Parameters
614
531
  ----------
615
- vars : Dict[str, str], default {}
616
- Dictionary of environment variables to set.
532
+ times : int, default 3
533
+ Number of times to retry this task.
534
+ minutes_between_retries : int, default 2
535
+ Number of minutes between retries.
617
536
  """
618
537
  ...
619
538
 
@@ -676,435 +595,345 @@ def timeout(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None],
676
595
  """
677
596
  ...
678
597
 
679
- def app_deploy(*, app_port: int, app_name: 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]]]:
680
- """
681
- Specifies that this step is used to deploy an instance of the app.
682
- Requires that self.app_name, self.app_port, self.entrypoint and self.deployDir is set.
683
-
684
-
685
- Parameters
686
- ----------
687
- app_port : int
688
- Number of GPUs to use.
689
- app_name : str
690
- Name of the app to deploy.
691
- """
692
- ...
693
-
694
598
  @typing.overload
695
- 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]]]:
599
+ 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]]]:
696
600
  """
697
- Specifies the PyPI packages for the step.
601
+ Specifies that the step will success under all circumstances.
698
602
 
699
- Information in this decorator will augment any
700
- attributes set in the `@pyi_base` flow-level decorator. Hence,
701
- you can use `@pypi_base` to set packages required by all
702
- steps and use `@pypi` to specify step-specific overrides.
603
+ The decorator will create an optional artifact, specified by `var`, which
604
+ contains the exception raised. You can use it to detect the presence
605
+ of errors, indicating that all happy-path artifacts produced by the step
606
+ are missing.
703
607
 
704
608
 
705
609
  Parameters
706
610
  ----------
707
- packages : Dict[str, str], default: {}
708
- Packages to use for this step. The key is the name of the package
709
- and the value is the version to use.
710
- python : str, optional, default: None
711
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
712
- that the version used will correspond to the version of the Python interpreter used to start the run.
611
+ var : str, optional, default None
612
+ Name of the artifact in which to store the caught exception.
613
+ If not specified, the exception is not stored.
614
+ print_exception : bool, default True
615
+ Determines whether or not the exception is printed to
616
+ stdout when caught.
713
617
  """
714
618
  ...
715
619
 
716
620
  @typing.overload
717
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
621
+ def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
718
622
  ...
719
623
 
720
624
  @typing.overload
721
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
625
+ def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
722
626
  ...
723
627
 
724
- 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):
628
+ 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):
725
629
  """
726
- Specifies the PyPI packages for the step.
630
+ Specifies that the step will success under all circumstances.
727
631
 
728
- Information in this decorator will augment any
729
- attributes set in the `@pyi_base` flow-level decorator. Hence,
730
- you can use `@pypi_base` to set packages required by all
731
- steps and use `@pypi` to specify step-specific overrides.
632
+ The decorator will create an optional artifact, specified by `var`, which
633
+ contains the exception raised. You can use it to detect the presence
634
+ of errors, indicating that all happy-path artifacts produced by the step
635
+ are missing.
732
636
 
733
637
 
734
638
  Parameters
735
639
  ----------
736
- packages : Dict[str, str], default: {}
737
- Packages to use for this step. The key is the name of the package
738
- and the value is the version to use.
739
- python : str, optional, default: None
740
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
741
- that the version used will correspond to the version of the Python interpreter used to start the run.
640
+ var : str, optional, default None
641
+ Name of the artifact in which to store the caught exception.
642
+ If not specified, the exception is not stored.
643
+ print_exception : bool, default True
644
+ Determines whether or not the exception is printed to
645
+ stdout when caught.
742
646
  """
743
647
  ...
744
648
 
745
649
  @typing.overload
746
- def resources(*, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
650
+ 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]]]:
747
651
  """
748
- Specifies the resources needed when executing this step.
749
-
750
- Use `@resources` to specify the resource requirements
751
- independently of the specific compute layer (`@batch`, `@kubernetes`).
652
+ Enables loading / saving of models within a step.
752
653
 
753
- You can choose the compute layer on the command line by executing e.g.
754
- ```
755
- python myflow.py run --with batch
756
- ```
757
- or
758
- ```
759
- python myflow.py run --with kubernetes
760
- ```
761
- which executes the flow on the desired system using the
762
- requirements specified in `@resources`.
763
654
 
764
655
 
765
656
  Parameters
766
657
  ----------
767
- cpu : int, default 1
768
- Number of CPUs required for this step.
769
- gpu : int, optional, default None
770
- Number of GPUs required for this step.
771
- disk : int, optional, default None
772
- Disk size (in MB) required for this step. Only applies on Kubernetes.
773
- memory : int, default 4096
774
- Memory size (in MB) required for this step.
775
- shared_memory : int, optional, default None
776
- The value for the size (in MiB) of the /dev/shm volume for this step.
777
- This parameter maps to the `--shm-size` option in Docker.
658
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
659
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
660
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
661
+ - `current.checkpoint`
662
+ - `current.model`
663
+ - `current.huggingface_hub`
664
+
665
+ 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
666
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
667
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
668
+
669
+ temp_dir_root : str, default: None
670
+ The root directory under which `current.model.loaded` will store loaded models
778
671
  """
779
672
  ...
780
673
 
781
674
  @typing.overload
782
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
675
+ def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
783
676
  ...
784
677
 
785
678
  @typing.overload
786
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
679
+ def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
787
680
  ...
788
681
 
789
- def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None):
682
+ 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):
790
683
  """
791
- Specifies the resources needed when executing this step.
792
-
793
- Use `@resources` to specify the resource requirements
794
- independently of the specific compute layer (`@batch`, `@kubernetes`).
684
+ Enables loading / saving of models within a step.
795
685
 
796
- You can choose the compute layer on the command line by executing e.g.
797
- ```
798
- python myflow.py run --with batch
799
- ```
800
- or
801
- ```
802
- python myflow.py run --with kubernetes
803
- ```
804
- which executes the flow on the desired system using the
805
- requirements specified in `@resources`.
806
686
 
807
687
 
808
688
  Parameters
809
689
  ----------
810
- cpu : int, default 1
811
- Number of CPUs required for this step.
812
- gpu : int, optional, default None
813
- Number of GPUs required for this step.
814
- disk : int, optional, default None
815
- Disk size (in MB) required for this step. Only applies on Kubernetes.
816
- memory : int, default 4096
817
- Memory size (in MB) required for this step.
818
- shared_memory : int, optional, default None
819
- The value for the size (in MiB) of the /dev/shm volume for this step.
820
- This parameter maps to the `--shm-size` option in Docker.
690
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
691
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
692
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
693
+ - `current.checkpoint`
694
+ - `current.model`
695
+ - `current.huggingface_hub`
696
+
697
+ 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
698
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
699
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
700
+
701
+ temp_dir_root : str, default: None
702
+ The root directory under which `current.model.loaded` will store loaded models
821
703
  """
822
704
  ...
823
705
 
824
- @typing.overload
825
- 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]]]:
706
+ def nvidia(*, gpu: int, gpu_type: str, queue_timeout: int) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
826
707
  """
827
- Specifies the Conda environment for the step.
828
-
829
- Information in this decorator will augment any
830
- attributes set in the `@conda_base` flow-level decorator. Hence,
831
- you can use `@conda_base` to set packages required by all
832
- steps and use `@conda` to specify step-specific overrides.
708
+ Specifies that this step should execute on DGX cloud.
833
709
 
834
710
 
835
711
  Parameters
836
712
  ----------
837
- packages : Dict[str, str], default {}
838
- Packages to use for this step. The key is the name of the package
839
- and the value is the version to use.
840
- libraries : Dict[str, str], default {}
841
- Supported for backward compatibility. When used with packages, packages will take precedence.
842
- python : str, optional, default None
843
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
844
- that the version used will correspond to the version of the Python interpreter used to start the run.
845
- disabled : bool, default False
846
- If set to True, disables @conda.
713
+ gpu : int
714
+ Number of GPUs to use.
715
+ gpu_type : str
716
+ Type of Nvidia GPU to use.
717
+ queue_timeout : int
718
+ Time to keep the job in NVCF's queue.
847
719
  """
848
720
  ...
849
721
 
850
- @typing.overload
851
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
852
- ...
853
-
854
- @typing.overload
855
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
856
- ...
857
-
858
- 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):
722
+ def huggingface_hub(*, temp_dir_root: typing.Optional[str] = None, load: typing.Union[typing.List[str], typing.List[typing.Tuple[typing.Dict, str]], typing.List[typing.Tuple[str, str]], typing.List[typing.Dict], 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]]]:
859
723
  """
860
- Specifies the Conda environment for the step.
861
-
862
- Information in this decorator will augment any
863
- attributes set in the `@conda_base` flow-level decorator. Hence,
864
- you can use `@conda_base` to set packages required by all
865
- steps and use `@conda` to specify step-specific overrides.
724
+ Decorator that helps cache, version and store models/datasets from huggingface hub.
866
725
 
867
726
 
868
727
  Parameters
869
728
  ----------
870
- packages : Dict[str, str], default {}
871
- Packages to use for this step. The key is the name of the package
872
- and the value is the version to use.
873
- libraries : Dict[str, str], default {}
874
- Supported for backward compatibility. When used with packages, packages will take precedence.
875
- python : str, optional, default None
876
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
877
- that the version used will correspond to the version of the Python interpreter used to start the run.
878
- disabled : bool, default False
879
- If set to True, disables @conda.
729
+ temp_dir_root : str, optional
730
+ The root directory that will hold the temporary directory where objects will be downloaded.
731
+
732
+ load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
733
+ The list of repos (models/datasets) to load.
734
+
735
+ Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
736
+
737
+ - If repo (model/dataset) is not found in the datastore:
738
+ - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
739
+ - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
740
+ - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
741
+
742
+ - If repo is found in the datastore:
743
+ - Loads it directly from datastore to local path (can be temporary directory or specified path)
880
744
  """
881
745
  ...
882
746
 
883
747
  @typing.overload
884
- 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]]]:
748
+ 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]]]:
885
749
  """
886
- Specifies secrets to be retrieved and injected as environment variables prior to
887
- the execution of a step.
750
+ Creates a human-readable report, a Metaflow Card, after this step completes.
751
+
752
+ Note that you may add multiple `@card` decorators in a step with different parameters.
888
753
 
889
754
 
890
755
  Parameters
891
756
  ----------
892
- sources : List[Union[str, Dict[str, Any]]], default: []
893
- List of secret specs, defining how the secrets are to be retrieved
757
+ type : str, default 'default'
758
+ Card type.
759
+ id : str, optional, default None
760
+ If multiple cards are present, use this id to identify this card.
761
+ options : Dict[str, Any], default {}
762
+ Options passed to the card. The contents depend on the card type.
763
+ timeout : int, default 45
764
+ Interrupt reporting if it takes more than this many seconds.
894
765
  """
895
766
  ...
896
767
 
897
768
  @typing.overload
898
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
769
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
899
770
  ...
900
771
 
901
772
  @typing.overload
902
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
773
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
903
774
  ...
904
775
 
905
- 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]]] = []):
776
+ 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):
906
777
  """
907
- Specifies secrets to be retrieved and injected as environment variables prior to
908
- the execution of a step.
778
+ Creates a human-readable report, a Metaflow Card, after this step completes.
779
+
780
+ Note that you may add multiple `@card` decorators in a step with different parameters.
909
781
 
910
782
 
911
783
  Parameters
912
784
  ----------
913
- sources : List[Union[str, Dict[str, Any]]], default: []
914
- List of secret specs, defining how the secrets are to be retrieved
785
+ type : str, default 'default'
786
+ Card type.
787
+ id : str, optional, default None
788
+ If multiple cards are present, use this id to identify this card.
789
+ options : Dict[str, Any], default {}
790
+ Options passed to the card. The contents depend on the card type.
791
+ timeout : int, default 45
792
+ Interrupt reporting if it takes more than this many seconds.
915
793
  """
916
794
  ...
917
795
 
918
796
  @typing.overload
919
- 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]]]:
797
+ 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]]]:
920
798
  """
921
- Specifies the number of times the task corresponding
922
- to a step needs to be retried.
923
-
924
- This decorator is useful for handling transient errors, such as networking issues.
925
- If your task contains operations that can't be retried safely, e.g. database updates,
926
- it is advisable to annotate it with `@retry(times=0)`.
799
+ Enables checkpointing for a step.
927
800
 
928
- This can be used in conjunction with the `@catch` decorator. The `@catch`
929
- decorator will execute a no-op task after all retries have been exhausted,
930
- ensuring that the flow execution can continue.
931
801
 
932
802
 
933
803
  Parameters
934
804
  ----------
935
- times : int, default 3
936
- Number of times to retry this task.
937
- minutes_between_retries : int, default 2
938
- Number of minutes between retries.
805
+ load_policy : str, default: "fresh"
806
+ The policy for loading the checkpoint. The following policies are supported:
807
+ - "eager": Loads the the latest available checkpoint within the namespace.
808
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
809
+ will be loaded at the start of the task.
810
+ - "none": Do not load any checkpoint
811
+ - "fresh": Loads the lastest checkpoint created within the running Task.
812
+ This mode helps loading checkpoints across various retry attempts of the same task.
813
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
814
+ created within the task will be loaded when the task is retries execution on failure.
815
+
816
+ temp_dir_root : str, default: None
817
+ The root directory under which `current.checkpoint.directory` will be created.
939
818
  """
940
819
  ...
941
820
 
942
821
  @typing.overload
943
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
822
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
944
823
  ...
945
824
 
946
825
  @typing.overload
947
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
826
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
948
827
  ...
949
828
 
950
- 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):
829
+ 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):
951
830
  """
952
- Specifies the number of times the task corresponding
953
- to a step needs to be retried.
831
+ Enables checkpointing for a step.
954
832
 
955
- This decorator is useful for handling transient errors, such as networking issues.
956
- If your task contains operations that can't be retried safely, e.g. database updates,
957
- it is advisable to annotate it with `@retry(times=0)`.
958
-
959
- This can be used in conjunction with the `@catch` decorator. The `@catch`
960
- decorator will execute a no-op task after all retries have been exhausted,
961
- ensuring that the flow execution can continue.
962
833
 
963
834
 
964
835
  Parameters
965
836
  ----------
966
- times : int, default 3
967
- Number of times to retry this task.
968
- minutes_between_retries : int, default 2
969
- Number of minutes between retries.
837
+ load_policy : str, default: "fresh"
838
+ The policy for loading the checkpoint. The following policies are supported:
839
+ - "eager": Loads the the latest available checkpoint within the namespace.
840
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
841
+ will be loaded at the start of the task.
842
+ - "none": Do not load any checkpoint
843
+ - "fresh": Loads the lastest checkpoint created within the running Task.
844
+ This mode helps loading checkpoints across various retry attempts of the same task.
845
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
846
+ created within the task will be loaded when the task is retries execution on failure.
847
+
848
+ temp_dir_root : str, default: None
849
+ The root directory under which `current.checkpoint.directory` will be created.
970
850
  """
971
851
  ...
972
852
 
973
- @typing.overload
974
- 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]]:
853
+ def nim(*, models: "list[NIM]", backend: str, queue_timeout: int) -> 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]]]:
975
854
  """
976
- Specifies the flow(s) that this flow depends on.
977
-
978
- ```
979
- @trigger_on_finish(flow='FooFlow')
980
- ```
981
- or
982
- ```
983
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
984
- ```
985
- This decorator respects the @project decorator and triggers the flow
986
- when upstream runs within the same namespace complete successfully
855
+ This decorator is used to run NIM containers in Metaflow tasks as sidecars.
987
856
 
988
- Additionally, you can specify project aware upstream flow dependencies
989
- by specifying the fully qualified project_flow_name.
990
- ```
991
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
992
- ```
993
- or
994
- ```
995
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
996
- ```
857
+ User code call
858
+ -----------
859
+ @nim(
860
+ models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
861
+ backend='managed'
862
+ )
997
863
 
998
- You can also specify just the project or project branch (other values will be
999
- inferred from the current project or project branch):
1000
- ```
1001
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1002
- ```
864
+ Valid backend options
865
+ ---------------------
866
+ - 'managed': Outerbounds selects a compute provider based on the model.
1003
867
 
1004
- Note that `branch` is typically one of:
1005
- - `prod`
1006
- - `user.bob`
1007
- - `test.my_experiment`
1008
- - `prod.staging`
868
+ Valid model options
869
+ ----------------
870
+ - 'meta/llama3-8b-instruct': 8B parameter model
871
+ - 'meta/llama3-70b-instruct': 70B parameter model
872
+ - any model here: https://nvcf.ngc.nvidia.com/functions?filter=nvidia-functions
1009
873
 
1010
874
 
1011
875
  Parameters
1012
876
  ----------
1013
- flow : Union[str, Dict[str, str]], optional, default None
1014
- Upstream flow dependency for this flow.
1015
- flows : List[Union[str, Dict[str, str]]], default []
1016
- Upstream flow dependencies for this flow.
1017
- options : Dict[str, Any], default {}
1018
- Backend-specific configuration for tuning eventing behavior.
877
+ models: list[NIM]
878
+ List of NIM containers running models in sidecars.
879
+ backend: str
880
+ Compute provider to run the NIM container.
881
+ queue_timeout : int
882
+ Time to keep the job in NVCF's queue.
1019
883
  """
1020
884
  ...
1021
885
 
1022
886
  @typing.overload
1023
- def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
887
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
888
+ """
889
+ Decorator prototype for all step decorators. This function gets specialized
890
+ and imported for all decorators types by _import_plugin_decorators().
891
+ """
1024
892
  ...
1025
893
 
1026
- def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}):
894
+ @typing.overload
895
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
896
+ ...
897
+
898
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1027
899
  """
1028
- Specifies the flow(s) that this flow depends on.
1029
-
1030
- ```
1031
- @trigger_on_finish(flow='FooFlow')
1032
- ```
1033
- or
1034
- ```
1035
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1036
- ```
1037
- This decorator respects the @project decorator and triggers the flow
1038
- when upstream runs within the same namespace complete successfully
1039
-
1040
- Additionally, you can specify project aware upstream flow dependencies
1041
- by specifying the fully qualified project_flow_name.
1042
- ```
1043
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1044
- ```
1045
- or
1046
- ```
1047
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1048
- ```
1049
-
1050
- You can also specify just the project or project branch (other values will be
1051
- inferred from the current project or project branch):
1052
- ```
1053
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1054
- ```
1055
-
1056
- Note that `branch` is typically one of:
1057
- - `prod`
1058
- - `user.bob`
1059
- - `test.my_experiment`
1060
- - `prod.staging`
1061
-
1062
-
1063
- Parameters
1064
- ----------
1065
- flow : Union[str, Dict[str, str]], optional, default None
1066
- Upstream flow dependency for this flow.
1067
- flows : List[Union[str, Dict[str, str]]], default []
1068
- Upstream flow dependencies for this flow.
1069
- options : Dict[str, Any], default {}
1070
- Backend-specific configuration for tuning eventing behavior.
900
+ Decorator prototype for all step decorators. This function gets specialized
901
+ and imported for all decorators types by _import_plugin_decorators().
1071
902
  """
1072
903
  ...
1073
904
 
1074
905
  @typing.overload
1075
- def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
906
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1076
907
  """
1077
- Specifies the PyPI packages for all steps of the flow.
1078
-
1079
- Use `@pypi_base` to set common packages required by all
1080
- steps and use `@pypi` to specify step-specific overrides.
1081
-
1082
- Parameters
1083
- ----------
1084
- packages : Dict[str, str], default: {}
1085
- Packages to use for this flow. The key is the name of the package
1086
- and the value is the version to use.
1087
- python : str, optional, default: None
1088
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1089
- that the version used will correspond to the version of the Python interpreter used to start the run.
908
+ Internal decorator to support Fast bakery
1090
909
  """
1091
910
  ...
1092
911
 
1093
912
  @typing.overload
1094
- def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
913
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1095
914
  ...
1096
915
 
1097
- def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
916
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1098
917
  """
1099
- Specifies the PyPI packages for all steps of the flow.
918
+ Internal decorator to support Fast bakery
919
+ """
920
+ ...
921
+
922
+ @typing.overload
923
+ 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]]]:
924
+ """
925
+ Specifies the PyPI packages for the step.
1100
926
 
1101
- Use `@pypi_base` to set common packages required by all
927
+ Information in this decorator will augment any
928
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
929
+ you can use `@pypi_base` to set packages required by all
1102
930
  steps and use `@pypi` to specify step-specific overrides.
1103
931
 
932
+
1104
933
  Parameters
1105
934
  ----------
1106
935
  packages : Dict[str, str], default: {}
1107
- Packages to use for this flow. The key is the name of the package
936
+ Packages to use for this step. The key is the name of the package
1108
937
  and the value is the version to use.
1109
938
  python : str, optional, default: None
1110
939
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
@@ -1112,203 +941,32 @@ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packag
1112
941
  """
1113
942
  ...
1114
943
 
1115
- 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]]:
1116
- """
1117
- 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)
1118
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1119
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1120
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
1121
- starts only after all sensors finish.
1122
-
1123
-
1124
- Parameters
1125
- ----------
1126
- timeout : int
1127
- Time, in seconds before the task times out and fails. (Default: 3600)
1128
- poke_interval : int
1129
- Time in seconds that the job should wait in between each try. (Default: 60)
1130
- mode : str
1131
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1132
- exponential_backoff : bool
1133
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1134
- pool : str
1135
- the slot pool this task should run in,
1136
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1137
- soft_fail : bool
1138
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1139
- name : str
1140
- Name of the sensor on Airflow
1141
- description : str
1142
- Description of sensor in the Airflow UI
1143
- bucket_key : Union[str, List[str]]
1144
- The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1145
- When it's specified as a full s3:// url, please leave `bucket_name` as None
1146
- bucket_name : str
1147
- Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1148
- When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1149
- wildcard_match : bool
1150
- whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1151
- aws_conn_id : str
1152
- a reference to the s3 connection on Airflow. (Default: None)
1153
- verify : bool
1154
- Whether or not to verify SSL certificates for S3 connection. (Default: None)
1155
- """
944
+ @typing.overload
945
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1156
946
  ...
1157
947
 
1158
- def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None):
1159
- """
1160
- Allows setting external datastores to save data for the
1161
- `@checkpoint`/`@model`/`@huggingface_hub` decorators.
1162
-
1163
- This decorator is useful when users wish to save data to a different datastore
1164
- than what is configured in Metaflow. This can be for variety of reasons:
1165
-
1166
- 1. Data security: The objects needs to be stored in a bucket (object storage) that is not accessible by other flows.
1167
- 2. Data Locality: The location where the task is executing is not located in the same region as the datastore.
1168
- - Example: Metaflow datastore lives in US East, but the task is executing in Finland datacenters.
1169
- 3. Data Lifecycle Policies: The objects need to be archived / managed separately from the Metaflow managed objects.
1170
- - Example: Flow is training very large models that need to be stored separately and will be deleted more aggressively than the Metaflow managed objects.
1171
-
1172
- Usage:
1173
- ----------
1174
-
1175
- - Using a custom IAM role to access the datastore.
1176
-
1177
- ```python
1178
- @with_artifact_store(
1179
- type="s3",
1180
- config=lambda: {
1181
- "root": "s3://my-bucket-foo/path/to/root",
1182
- "role_arn": ROLE,
1183
- },
1184
- )
1185
- class MyFlow(FlowSpec):
1186
-
1187
- @checkpoint
1188
- @step
1189
- def start(self):
1190
- with open("my_file.txt", "w") as f:
1191
- f.write("Hello, World!")
1192
- self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1193
- self.next(self.end)
1194
-
1195
- ```
1196
-
1197
- - Using credentials to access the s3-compatible datastore.
1198
-
1199
- ```python
1200
- @with_artifact_store(
1201
- type="s3",
1202
- config=lambda: {
1203
- "root": "s3://my-bucket-foo/path/to/root",
1204
- "client_params": {
1205
- "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1206
- "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1207
- },
1208
- },
1209
- )
1210
- class MyFlow(FlowSpec):
1211
-
1212
- @checkpoint
1213
- @step
1214
- def start(self):
1215
- with open("my_file.txt", "w") as f:
1216
- f.write("Hello, World!")
1217
- self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1218
- self.next(self.end)
1219
-
1220
- ```
1221
-
1222
- - Accessing objects stored in external datastores after task execution.
1223
-
1224
- ```python
1225
- run = Run("CheckpointsTestsFlow/8992")
1226
- with artifact_store_from(run=run, config={
1227
- "client_params": {
1228
- "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1229
- "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1230
- },
1231
- }):
1232
- with Checkpoint() as cp:
1233
- latest = cp.list(
1234
- task=run["start"].task
1235
- )[0]
1236
- print(latest)
1237
- cp.load(
1238
- latest,
1239
- "test-checkpoints"
1240
- )
1241
-
1242
- task = Task("TorchTuneFlow/8484/train/53673")
1243
- with artifact_store_from(run=run, config={
1244
- "client_params": {
1245
- "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1246
- "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1247
- },
1248
- }):
1249
- load_model(
1250
- task.data.model_ref,
1251
- "test-models"
1252
- )
1253
- ```
1254
- Parameters:
1255
- ----------
1256
-
1257
- type: str
1258
- The type of the datastore. Can be one of 's3', 'gcs', 'azure' or any other supported metaflow Datastore.
1259
-
1260
- config: dict or Callable
1261
- Dictionary of configuration options for the datastore. The following keys are required:
1262
- - root: The root path in the datastore where the data will be saved. (needs to be in the format expected by the datastore)
1263
- - example: 's3://bucket-name/path/to/root'
1264
- - example: 'gs://bucket-name/path/to/root'
1265
- - example: 'https://myblockacc.blob.core.windows.net/metaflow/'
1266
- - role_arn (optional): AWS IAM role to access s3 bucket (only when `type` is 's3')
1267
- - session_vars (optional): AWS session variables to access s3 bucket (only when `type` is 's3')
1268
- - client_params (optional): AWS client parameters to access s3 bucket (only when `type` is 's3')
1269
- """
948
+ @typing.overload
949
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1270
950
  ...
1271
951
 
1272
- 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]]:
952
+ 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):
1273
953
  """
1274
- 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.
1275
- 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.
954
+ Specifies the PyPI packages for the step.
955
+
956
+ Information in this decorator will augment any
957
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
958
+ you can use `@pypi_base` to set packages required by all
959
+ steps and use `@pypi` to specify step-specific overrides.
1276
960
 
1277
961
 
1278
962
  Parameters
1279
963
  ----------
1280
- timeout : int
1281
- Time, in seconds before the task times out and fails. (Default: 3600)
1282
- poke_interval : int
1283
- Time in seconds that the job should wait in between each try. (Default: 60)
1284
- mode : str
1285
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1286
- exponential_backoff : bool
1287
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1288
- pool : str
1289
- the slot pool this task should run in,
1290
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1291
- soft_fail : bool
1292
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1293
- name : str
1294
- Name of the sensor on Airflow
1295
- description : str
1296
- Description of sensor in the Airflow UI
1297
- external_dag_id : str
1298
- The dag_id that contains the task you want to wait for.
1299
- external_task_ids : List[str]
1300
- The list of task_ids that you want to wait for.
1301
- If None (default value) the sensor waits for the DAG. (Default: None)
1302
- allowed_states : List[str]
1303
- Iterable of allowed states, (Default: ['success'])
1304
- failed_states : List[str]
1305
- Iterable of failed or dis-allowed states. (Default: None)
1306
- execution_delta : datetime.timedelta
1307
- time difference with the previous execution to look at,
1308
- the default is the same logical date as the current task or DAG. (Default: None)
1309
- check_existence: bool
1310
- Set to True to check if the external task exists or check if
1311
- the DAG to wait for exists. (Default: True)
964
+ packages : Dict[str, str], default: {}
965
+ Packages to use for this step. The key is the name of the package
966
+ and the value is the version to use.
967
+ python : str, optional, default: None
968
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
969
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1312
970
  """
1313
971
  ...
1314
972
 
@@ -1434,58 +1092,202 @@ def trigger(*, event: typing.Union[str, typing.Dict[str, typing.Any], None] = No
1434
1092
 
1435
1093
  Parameters
1436
1094
  ----------
1437
- event : Union[str, Dict[str, Any]], optional, default None
1438
- Event dependency for this flow.
1439
- events : List[Union[str, Dict[str, Any]]], default []
1440
- Events dependency for this flow.
1095
+ event : Union[str, Dict[str, Any]], optional, default None
1096
+ Event dependency for this flow.
1097
+ events : List[Union[str, Dict[str, Any]]], default []
1098
+ Events dependency for this flow.
1099
+ options : Dict[str, Any], default {}
1100
+ Backend-specific configuration for tuning eventing behavior.
1101
+ """
1102
+ ...
1103
+
1104
+ @typing.overload
1105
+ def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1106
+ ...
1107
+
1108
+ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: typing.Union[str, typing.Dict[str, typing.Any], None] = None, events: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], options: typing.Dict[str, typing.Any] = {}):
1109
+ """
1110
+ Specifies the event(s) that this flow depends on.
1111
+
1112
+ ```
1113
+ @trigger(event='foo')
1114
+ ```
1115
+ or
1116
+ ```
1117
+ @trigger(events=['foo', 'bar'])
1118
+ ```
1119
+
1120
+ Additionally, you can specify the parameter mappings
1121
+ to map event payload to Metaflow parameters for the flow.
1122
+ ```
1123
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1124
+ ```
1125
+ or
1126
+ ```
1127
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1128
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1129
+ ```
1130
+
1131
+ 'parameters' can also be a list of strings and tuples like so:
1132
+ ```
1133
+ @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1134
+ ```
1135
+ This is equivalent to:
1136
+ ```
1137
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1138
+ ```
1139
+
1140
+
1141
+ Parameters
1142
+ ----------
1143
+ event : Union[str, Dict[str, Any]], optional, default None
1144
+ Event dependency for this flow.
1145
+ events : List[Union[str, Dict[str, Any]]], default []
1146
+ Events dependency for this flow.
1147
+ options : Dict[str, Any], default {}
1148
+ Backend-specific configuration for tuning eventing behavior.
1149
+ """
1150
+ ...
1151
+
1152
+ 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]]:
1153
+ """
1154
+ 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)
1155
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1156
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1157
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
1158
+ starts only after all sensors finish.
1159
+
1160
+
1161
+ Parameters
1162
+ ----------
1163
+ timeout : int
1164
+ Time, in seconds before the task times out and fails. (Default: 3600)
1165
+ poke_interval : int
1166
+ Time in seconds that the job should wait in between each try. (Default: 60)
1167
+ mode : str
1168
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1169
+ exponential_backoff : bool
1170
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1171
+ pool : str
1172
+ the slot pool this task should run in,
1173
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1174
+ soft_fail : bool
1175
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1176
+ name : str
1177
+ Name of the sensor on Airflow
1178
+ description : str
1179
+ Description of sensor in the Airflow UI
1180
+ bucket_key : Union[str, List[str]]
1181
+ The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1182
+ When it's specified as a full s3:// url, please leave `bucket_name` as None
1183
+ bucket_name : str
1184
+ Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1185
+ When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1186
+ wildcard_match : bool
1187
+ whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1188
+ aws_conn_id : str
1189
+ a reference to the s3 connection on Airflow. (Default: None)
1190
+ verify : bool
1191
+ Whether or not to verify SSL certificates for S3 connection. (Default: None)
1192
+ """
1193
+ ...
1194
+
1195
+ @typing.overload
1196
+ 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]]:
1197
+ """
1198
+ Specifies the flow(s) that this flow depends on.
1199
+
1200
+ ```
1201
+ @trigger_on_finish(flow='FooFlow')
1202
+ ```
1203
+ or
1204
+ ```
1205
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1206
+ ```
1207
+ This decorator respects the @project decorator and triggers the flow
1208
+ when upstream runs within the same namespace complete successfully
1209
+
1210
+ Additionally, you can specify project aware upstream flow dependencies
1211
+ by specifying the fully qualified project_flow_name.
1212
+ ```
1213
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1214
+ ```
1215
+ or
1216
+ ```
1217
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1218
+ ```
1219
+
1220
+ You can also specify just the project or project branch (other values will be
1221
+ inferred from the current project or project branch):
1222
+ ```
1223
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1224
+ ```
1225
+
1226
+ Note that `branch` is typically one of:
1227
+ - `prod`
1228
+ - `user.bob`
1229
+ - `test.my_experiment`
1230
+ - `prod.staging`
1231
+
1232
+
1233
+ Parameters
1234
+ ----------
1235
+ flow : Union[str, Dict[str, str]], optional, default None
1236
+ Upstream flow dependency for this flow.
1237
+ flows : List[Union[str, Dict[str, str]]], default []
1238
+ Upstream flow dependencies for this flow.
1441
1239
  options : Dict[str, Any], default {}
1442
1240
  Backend-specific configuration for tuning eventing behavior.
1443
1241
  """
1444
1242
  ...
1445
1243
 
1446
1244
  @typing.overload
1447
- def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1245
+ def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1448
1246
  ...
1449
1247
 
1450
- def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: typing.Union[str, typing.Dict[str, typing.Any], None] = None, events: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], options: typing.Dict[str, typing.Any] = {}):
1248
+ def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}):
1451
1249
  """
1452
- Specifies the event(s) that this flow depends on.
1250
+ Specifies the flow(s) that this flow depends on.
1453
1251
 
1454
1252
  ```
1455
- @trigger(event='foo')
1253
+ @trigger_on_finish(flow='FooFlow')
1456
1254
  ```
1457
1255
  or
1458
1256
  ```
1459
- @trigger(events=['foo', 'bar'])
1257
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1460
1258
  ```
1259
+ This decorator respects the @project decorator and triggers the flow
1260
+ when upstream runs within the same namespace complete successfully
1461
1261
 
1462
- Additionally, you can specify the parameter mappings
1463
- to map event payload to Metaflow parameters for the flow.
1262
+ Additionally, you can specify project aware upstream flow dependencies
1263
+ by specifying the fully qualified project_flow_name.
1464
1264
  ```
1465
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1265
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1466
1266
  ```
1467
1267
  or
1468
1268
  ```
1469
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1470
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1269
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1471
1270
  ```
1472
1271
 
1473
- 'parameters' can also be a list of strings and tuples like so:
1474
- ```
1475
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1476
- ```
1477
- This is equivalent to:
1272
+ You can also specify just the project or project branch (other values will be
1273
+ inferred from the current project or project branch):
1478
1274
  ```
1479
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1275
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1480
1276
  ```
1481
1277
 
1278
+ Note that `branch` is typically one of:
1279
+ - `prod`
1280
+ - `user.bob`
1281
+ - `test.my_experiment`
1282
+ - `prod.staging`
1283
+
1482
1284
 
1483
1285
  Parameters
1484
1286
  ----------
1485
- event : Union[str, Dict[str, Any]], optional, default None
1486
- Event dependency for this flow.
1487
- events : List[Union[str, Dict[str, Any]]], default []
1488
- Events dependency for this flow.
1287
+ flow : Union[str, Dict[str, str]], optional, default None
1288
+ Upstream flow dependency for this flow.
1289
+ flows : List[Union[str, Dict[str, str]]], default []
1290
+ Upstream flow dependencies for this flow.
1489
1291
  options : Dict[str, Any], default {}
1490
1292
  Backend-specific configuration for tuning eventing behavior.
1491
1293
  """
@@ -1542,5 +1344,203 @@ def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packa
1542
1344
  """
1543
1345
  ...
1544
1346
 
1347
+ @typing.overload
1348
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1349
+ """
1350
+ Specifies the PyPI packages for all steps of the flow.
1351
+
1352
+ Use `@pypi_base` to set common packages required by all
1353
+ steps and use `@pypi` to specify step-specific overrides.
1354
+
1355
+ Parameters
1356
+ ----------
1357
+ packages : Dict[str, str], default: {}
1358
+ Packages to use for this flow. The key is the name of the package
1359
+ and the value is the version to use.
1360
+ python : str, optional, default: None
1361
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1362
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1363
+ """
1364
+ ...
1365
+
1366
+ @typing.overload
1367
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1368
+ ...
1369
+
1370
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1371
+ """
1372
+ Specifies the PyPI packages for all steps of the flow.
1373
+
1374
+ Use `@pypi_base` to set common packages required by all
1375
+ steps and use `@pypi` to specify step-specific overrides.
1376
+
1377
+ Parameters
1378
+ ----------
1379
+ packages : Dict[str, str], default: {}
1380
+ Packages to use for this flow. The key is the name of the package
1381
+ and the value is the version to use.
1382
+ python : str, optional, default: None
1383
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1384
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1385
+ """
1386
+ ...
1387
+
1388
+ 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]]:
1389
+ """
1390
+ 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.
1391
+ 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.
1392
+
1393
+
1394
+ Parameters
1395
+ ----------
1396
+ timeout : int
1397
+ Time, in seconds before the task times out and fails. (Default: 3600)
1398
+ poke_interval : int
1399
+ Time in seconds that the job should wait in between each try. (Default: 60)
1400
+ mode : str
1401
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1402
+ exponential_backoff : bool
1403
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1404
+ pool : str
1405
+ the slot pool this task should run in,
1406
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1407
+ soft_fail : bool
1408
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1409
+ name : str
1410
+ Name of the sensor on Airflow
1411
+ description : str
1412
+ Description of sensor in the Airflow UI
1413
+ external_dag_id : str
1414
+ The dag_id that contains the task you want to wait for.
1415
+ external_task_ids : List[str]
1416
+ The list of task_ids that you want to wait for.
1417
+ If None (default value) the sensor waits for the DAG. (Default: None)
1418
+ allowed_states : List[str]
1419
+ Iterable of allowed states, (Default: ['success'])
1420
+ failed_states : List[str]
1421
+ Iterable of failed or dis-allowed states. (Default: None)
1422
+ execution_delta : datetime.timedelta
1423
+ time difference with the previous execution to look at,
1424
+ the default is the same logical date as the current task or DAG. (Default: None)
1425
+ check_existence: bool
1426
+ Set to True to check if the external task exists or check if
1427
+ the DAG to wait for exists. (Default: True)
1428
+ """
1429
+ ...
1430
+
1431
+ def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None):
1432
+ """
1433
+ Allows setting external datastores to save data for the
1434
+ `@checkpoint`/`@model`/`@huggingface_hub` decorators.
1435
+
1436
+ This decorator is useful when users wish to save data to a different datastore
1437
+ than what is configured in Metaflow. This can be for variety of reasons:
1438
+
1439
+ 1. Data security: The objects needs to be stored in a bucket (object storage) that is not accessible by other flows.
1440
+ 2. Data Locality: The location where the task is executing is not located in the same region as the datastore.
1441
+ - Example: Metaflow datastore lives in US East, but the task is executing in Finland datacenters.
1442
+ 3. Data Lifecycle Policies: The objects need to be archived / managed separately from the Metaflow managed objects.
1443
+ - Example: Flow is training very large models that need to be stored separately and will be deleted more aggressively than the Metaflow managed objects.
1444
+
1445
+ Usage:
1446
+ ----------
1447
+
1448
+ - Using a custom IAM role to access the datastore.
1449
+
1450
+ ```python
1451
+ @with_artifact_store(
1452
+ type="s3",
1453
+ config=lambda: {
1454
+ "root": "s3://my-bucket-foo/path/to/root",
1455
+ "role_arn": ROLE,
1456
+ },
1457
+ )
1458
+ class MyFlow(FlowSpec):
1459
+
1460
+ @checkpoint
1461
+ @step
1462
+ def start(self):
1463
+ with open("my_file.txt", "w") as f:
1464
+ f.write("Hello, World!")
1465
+ self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1466
+ self.next(self.end)
1467
+
1468
+ ```
1469
+
1470
+ - Using credentials to access the s3-compatible datastore.
1471
+
1472
+ ```python
1473
+ @with_artifact_store(
1474
+ type="s3",
1475
+ config=lambda: {
1476
+ "root": "s3://my-bucket-foo/path/to/root",
1477
+ "client_params": {
1478
+ "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1479
+ "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1480
+ },
1481
+ },
1482
+ )
1483
+ class MyFlow(FlowSpec):
1484
+
1485
+ @checkpoint
1486
+ @step
1487
+ def start(self):
1488
+ with open("my_file.txt", "w") as f:
1489
+ f.write("Hello, World!")
1490
+ self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1491
+ self.next(self.end)
1492
+
1493
+ ```
1494
+
1495
+ - Accessing objects stored in external datastores after task execution.
1496
+
1497
+ ```python
1498
+ run = Run("CheckpointsTestsFlow/8992")
1499
+ with artifact_store_from(run=run, config={
1500
+ "client_params": {
1501
+ "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1502
+ "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1503
+ },
1504
+ }):
1505
+ with Checkpoint() as cp:
1506
+ latest = cp.list(
1507
+ task=run["start"].task
1508
+ )[0]
1509
+ print(latest)
1510
+ cp.load(
1511
+ latest,
1512
+ "test-checkpoints"
1513
+ )
1514
+
1515
+ task = Task("TorchTuneFlow/8484/train/53673")
1516
+ with artifact_store_from(run=run, config={
1517
+ "client_params": {
1518
+ "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1519
+ "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1520
+ },
1521
+ }):
1522
+ load_model(
1523
+ task.data.model_ref,
1524
+ "test-models"
1525
+ )
1526
+ ```
1527
+ Parameters:
1528
+ ----------
1529
+
1530
+ type: str
1531
+ The type of the datastore. Can be one of 's3', 'gcs', 'azure' or any other supported metaflow Datastore.
1532
+
1533
+ config: dict or Callable
1534
+ Dictionary of configuration options for the datastore. The following keys are required:
1535
+ - root: The root path in the datastore where the data will be saved. (needs to be in the format expected by the datastore)
1536
+ - example: 's3://bucket-name/path/to/root'
1537
+ - example: 'gs://bucket-name/path/to/root'
1538
+ - example: 'https://myblockacc.blob.core.windows.net/metaflow/'
1539
+ - role_arn (optional): AWS IAM role to access s3 bucket (only when `type` is 's3')
1540
+ - session_vars (optional): AWS session variables to access s3 bucket (only when `type` is 's3')
1541
+ - client_params (optional): AWS client parameters to access s3 bucket (only when `type` is 's3')
1542
+ """
1543
+ ...
1544
+
1545
1545
  pkg_name: str
1546
1546