ob-metaflow-stubs 6.0.7.4__py2.py3-none-any.whl → 6.0.8.1__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.

Potentially problematic release.


This version of ob-metaflow-stubs might be problematic. Click here for more details.

Files changed (262) hide show
  1. metaflow-stubs/__init__.pyi +1008 -1008
  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 +5 -5
  8. metaflow-stubs/client/filecache.pyi +1 -1
  9. metaflow-stubs/events.pyi +2 -2
  10. metaflow-stubs/exception.pyi +1 -1
  11. metaflow-stubs/flowspec.pyi +4 -4
  12. metaflow-stubs/generated_for.txt +1 -1
  13. metaflow-stubs/includefile.pyi +3 -3
  14. metaflow-stubs/meta_files.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 +2 -2
  18. metaflow-stubs/metadata_provider/util.pyi +1 -1
  19. metaflow-stubs/metaflow_config.pyi +1 -1
  20. metaflow-stubs/metaflow_current.pyi +70 -70
  21. metaflow-stubs/metaflow_git.pyi +1 -1
  22. metaflow-stubs/mf_extensions/__init__.pyi +1 -1
  23. metaflow-stubs/mf_extensions/obcheckpoint/__init__.pyi +1 -1
  24. metaflow-stubs/mf_extensions/obcheckpoint/plugins/__init__.pyi +1 -1
  25. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/__init__.pyi +1 -1
  26. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/__init__.pyi +1 -1
  27. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/async_cards.pyi +2 -2
  28. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/deco_injection_mixin.pyi +1 -1
  29. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/extra_components.pyi +3 -3
  30. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/__init__.pyi +1 -1
  31. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/__init__.pyi +1 -1
  32. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/checkpoint_lister.pyi +5 -5
  33. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/lineage_card.pyi +1 -1
  34. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/checkpoint_storage.pyi +3 -3
  35. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/constructors.pyi +1 -1
  36. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/core.pyi +3 -3
  37. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/decorator.pyi +3 -3
  38. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/exceptions.pyi +1 -1
  39. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/final_api.pyi +2 -2
  40. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/lineage.pyi +1 -1
  41. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/__init__.pyi +1 -1
  42. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/context.pyi +2 -2
  43. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/core.pyi +1 -1
  44. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/decorator.pyi +1 -1
  45. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/exceptions.pyi +1 -1
  46. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/task_utils.pyi +2 -2
  47. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/utils.pyi +1 -1
  48. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastructures.pyi +2 -2
  49. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/exceptions.pyi +1 -1
  50. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/__init__.pyi +1 -1
  51. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/decorator.pyi +2 -2
  52. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/__init__.pyi +1 -1
  53. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/core.pyi +2 -2
  54. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/exceptions.pyi +1 -1
  55. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/model_storage.pyi +3 -3
  56. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/__init__.pyi +1 -1
  57. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/flowspec_utils.pyi +1 -1
  58. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/general.pyi +1 -1
  59. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/identity_utils.pyi +1 -1
  60. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/__init__.pyi +1 -1
  61. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/base.pyi +1 -1
  62. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/tar.pyi +2 -2
  63. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/tar_utils.pyi +2 -2
  64. metaflow-stubs/mf_extensions/outerbounds/__init__.pyi +1 -1
  65. metaflow-stubs/mf_extensions/outerbounds/plugins/__init__.pyi +1 -1
  66. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/__init__.pyi +1 -1
  67. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/__init__.pyi +1 -1
  68. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/_state_machine.pyi +1 -1
  69. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/_vendor/__init__.pyi +1 -1
  70. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/_vendor/spinner/__init__.pyi +1 -1
  71. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/_vendor/spinner/spinners.pyi +1 -1
  72. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/app_cli.pyi +2 -2
  73. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/app_config.pyi +1 -1
  74. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/capsule.pyi +3 -3
  75. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/click_importer.pyi +1 -1
  76. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/code_package/__init__.pyi +1 -1
  77. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/code_package/code_packager.pyi +1 -1
  78. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/__init__.pyi +1 -1
  79. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/cli_generator.pyi +1 -1
  80. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/config_utils.pyi +3 -3
  81. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/schema_export.pyi +1 -1
  82. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/typed_configs.pyi +2 -2
  83. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/unified_config.pyi +2 -2
  84. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/dependencies.pyi +2 -2
  85. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/deployer.pyi +4 -4
  86. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/experimental/__init__.pyi +1 -1
  87. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/perimeters.pyi +1 -1
  88. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/utils.pyi +2 -2
  89. metaflow-stubs/mf_extensions/outerbounds/plugins/aws/__init__.pyi +1 -1
  90. metaflow-stubs/mf_extensions/outerbounds/plugins/aws/assume_role_decorator.pyi +1 -1
  91. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/__init__.pyi +1 -1
  92. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/async_cards.pyi +2 -2
  93. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/injector.pyi +1 -1
  94. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/__init__.pyi +1 -1
  95. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/coreweave.pyi +2 -2
  96. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/nebius.pyi +2 -2
  97. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/__init__.pyi +1 -1
  98. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/baker.pyi +3 -3
  99. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/docker_environment.pyi +2 -2
  100. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/fast_bakery.pyi +1 -1
  101. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/__init__.pyi +1 -1
  102. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/pod_killer.pyi +1 -1
  103. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/__init__.pyi +1 -1
  104. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/constants.pyi +1 -1
  105. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/exceptions.pyi +1 -1
  106. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/ollama.pyi +1 -1
  107. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/status_card.pyi +1 -1
  108. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/__init__.pyi +1 -1
  109. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/snowflake.pyi +1 -1
  110. metaflow-stubs/mf_extensions/outerbounds/profilers/__init__.pyi +1 -1
  111. metaflow-stubs/mf_extensions/outerbounds/profilers/gpu.pyi +1 -1
  112. metaflow-stubs/mf_extensions/outerbounds/remote_config.pyi +2 -2
  113. metaflow-stubs/mf_extensions/outerbounds/toplevel/__init__.pyi +1 -1
  114. metaflow-stubs/mf_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.pyi +1 -1
  115. metaflow-stubs/mf_extensions/outerbounds/toplevel/s3_proxy.pyi +1 -1
  116. metaflow-stubs/multicore_utils.pyi +1 -1
  117. metaflow-stubs/ob_internal.pyi +1 -1
  118. metaflow-stubs/packaging_sys/__init__.pyi +5 -5
  119. metaflow-stubs/packaging_sys/backend.pyi +2 -2
  120. metaflow-stubs/packaging_sys/distribution_support.pyi +2 -2
  121. metaflow-stubs/packaging_sys/tar_backend.pyi +5 -5
  122. metaflow-stubs/packaging_sys/utils.pyi +1 -1
  123. metaflow-stubs/packaging_sys/v1.pyi +2 -2
  124. metaflow-stubs/parameters.pyi +3 -3
  125. metaflow-stubs/plugins/__init__.pyi +12 -12
  126. metaflow-stubs/plugins/airflow/__init__.pyi +1 -1
  127. metaflow-stubs/plugins/airflow/airflow_utils.pyi +1 -1
  128. metaflow-stubs/plugins/airflow/exception.pyi +1 -1
  129. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +1 -1
  130. metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi +1 -1
  131. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +1 -1
  132. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +1 -1
  133. metaflow-stubs/plugins/argo/__init__.pyi +1 -1
  134. metaflow-stubs/plugins/argo/argo_client.pyi +1 -1
  135. metaflow-stubs/plugins/argo/argo_events.pyi +1 -1
  136. metaflow-stubs/plugins/argo/argo_workflows.pyi +2 -2
  137. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +2 -2
  138. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +2 -2
  139. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +2 -2
  140. metaflow-stubs/plugins/argo/exit_hooks.pyi +2 -2
  141. metaflow-stubs/plugins/aws/__init__.pyi +1 -1
  142. metaflow-stubs/plugins/aws/aws_client.pyi +1 -1
  143. metaflow-stubs/plugins/aws/aws_utils.pyi +1 -1
  144. metaflow-stubs/plugins/aws/batch/__init__.pyi +1 -1
  145. metaflow-stubs/plugins/aws/batch/batch.pyi +1 -1
  146. metaflow-stubs/plugins/aws/batch/batch_client.pyi +1 -1
  147. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +1 -1
  148. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +1 -1
  149. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +3 -3
  150. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +1 -1
  151. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +1 -1
  152. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +1 -1
  153. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +1 -1
  154. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +1 -1
  155. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer.pyi +2 -2
  156. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer_objects.pyi +2 -2
  157. metaflow-stubs/plugins/azure/__init__.pyi +1 -1
  158. metaflow-stubs/plugins/azure/azure_credential.pyi +1 -1
  159. metaflow-stubs/plugins/azure/azure_exceptions.pyi +1 -1
  160. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +3 -3
  161. metaflow-stubs/plugins/azure/azure_utils.pyi +1 -1
  162. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +1 -1
  163. metaflow-stubs/plugins/azure/includefile_support.pyi +1 -1
  164. metaflow-stubs/plugins/cards/__init__.pyi +1 -1
  165. metaflow-stubs/plugins/cards/card_client.pyi +2 -2
  166. metaflow-stubs/plugins/cards/card_creator.pyi +1 -1
  167. metaflow-stubs/plugins/cards/card_datastore.pyi +1 -1
  168. metaflow-stubs/plugins/cards/card_decorator.pyi +2 -2
  169. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +1 -1
  170. metaflow-stubs/plugins/cards/card_modules/basic.pyi +1 -1
  171. metaflow-stubs/plugins/cards/card_modules/card.pyi +1 -1
  172. metaflow-stubs/plugins/cards/card_modules/components.pyi +4 -4
  173. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +1 -1
  174. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +1 -1
  175. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +1 -1
  176. metaflow-stubs/plugins/cards/card_resolver.pyi +1 -1
  177. metaflow-stubs/plugins/cards/component_serializer.pyi +1 -1
  178. metaflow-stubs/plugins/cards/exception.pyi +1 -1
  179. metaflow-stubs/plugins/catch_decorator.pyi +2 -2
  180. metaflow-stubs/plugins/datatools/__init__.pyi +1 -1
  181. metaflow-stubs/plugins/datatools/local.pyi +1 -1
  182. metaflow-stubs/plugins/datatools/s3/__init__.pyi +1 -1
  183. metaflow-stubs/plugins/datatools/s3/s3.pyi +3 -3
  184. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +1 -1
  185. metaflow-stubs/plugins/datatools/s3/s3util.pyi +1 -1
  186. metaflow-stubs/plugins/debug_logger.pyi +1 -1
  187. metaflow-stubs/plugins/debug_monitor.pyi +1 -1
  188. metaflow-stubs/plugins/environment_decorator.pyi +1 -1
  189. metaflow-stubs/plugins/events_decorator.pyi +1 -1
  190. metaflow-stubs/plugins/exit_hook/__init__.pyi +1 -1
  191. metaflow-stubs/plugins/exit_hook/exit_hook_decorator.pyi +1 -1
  192. metaflow-stubs/plugins/frameworks/__init__.pyi +1 -1
  193. metaflow-stubs/plugins/frameworks/pytorch.pyi +1 -1
  194. metaflow-stubs/plugins/gcp/__init__.pyi +1 -1
  195. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +3 -3
  196. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +1 -1
  197. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +1 -1
  198. metaflow-stubs/plugins/gcp/gs_utils.pyi +1 -1
  199. metaflow-stubs/plugins/gcp/includefile_support.pyi +1 -1
  200. metaflow-stubs/plugins/kubernetes/__init__.pyi +1 -1
  201. metaflow-stubs/plugins/kubernetes/kube_utils.pyi +2 -2
  202. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +1 -1
  203. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +1 -1
  204. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +1 -1
  205. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +1 -1
  206. metaflow-stubs/plugins/kubernetes/spot_monitor_sidecar.pyi +1 -1
  207. metaflow-stubs/plugins/ollama/__init__.pyi +2 -2
  208. metaflow-stubs/plugins/optuna/__init__.pyi +1 -1
  209. metaflow-stubs/plugins/parallel_decorator.pyi +1 -1
  210. metaflow-stubs/plugins/perimeters.pyi +1 -1
  211. metaflow-stubs/plugins/project_decorator.pyi +1 -1
  212. metaflow-stubs/plugins/pypi/__init__.pyi +1 -1
  213. metaflow-stubs/plugins/pypi/conda_decorator.pyi +1 -1
  214. metaflow-stubs/plugins/pypi/conda_environment.pyi +3 -3
  215. metaflow-stubs/plugins/pypi/parsers.pyi +1 -1
  216. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +1 -1
  217. metaflow-stubs/plugins/pypi/pypi_environment.pyi +1 -1
  218. metaflow-stubs/plugins/pypi/utils.pyi +1 -1
  219. metaflow-stubs/plugins/resources_decorator.pyi +1 -1
  220. metaflow-stubs/plugins/retry_decorator.pyi +1 -1
  221. metaflow-stubs/plugins/secrets/__init__.pyi +1 -1
  222. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +2 -2
  223. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +1 -1
  224. metaflow-stubs/plugins/secrets/secrets_func.pyi +1 -1
  225. metaflow-stubs/plugins/secrets/secrets_spec.pyi +1 -1
  226. metaflow-stubs/plugins/secrets/utils.pyi +1 -1
  227. metaflow-stubs/plugins/snowflake/__init__.pyi +1 -1
  228. metaflow-stubs/plugins/storage_executor.pyi +1 -1
  229. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +2 -2
  230. metaflow-stubs/plugins/timeout_decorator.pyi +2 -2
  231. metaflow-stubs/plugins/torchtune/__init__.pyi +1 -1
  232. metaflow-stubs/plugins/uv/__init__.pyi +1 -1
  233. metaflow-stubs/plugins/uv/uv_environment.pyi +2 -2
  234. metaflow-stubs/profilers/__init__.pyi +1 -1
  235. metaflow-stubs/pylint_wrapper.pyi +1 -1
  236. metaflow-stubs/runner/__init__.pyi +1 -1
  237. metaflow-stubs/runner/deployer.pyi +5 -5
  238. metaflow-stubs/runner/deployer_impl.pyi +2 -2
  239. metaflow-stubs/runner/metaflow_runner.pyi +3 -3
  240. metaflow-stubs/runner/nbdeploy.pyi +1 -1
  241. metaflow-stubs/runner/nbrun.pyi +1 -1
  242. metaflow-stubs/runner/subprocess_manager.pyi +1 -1
  243. metaflow-stubs/runner/utils.pyi +3 -3
  244. metaflow-stubs/system/__init__.pyi +1 -1
  245. metaflow-stubs/system/system_logger.pyi +2 -2
  246. metaflow-stubs/system/system_monitor.pyi +1 -1
  247. metaflow-stubs/tagging_util.pyi +1 -1
  248. metaflow-stubs/tuple_util.pyi +1 -1
  249. metaflow-stubs/user_configs/__init__.pyi +1 -1
  250. metaflow-stubs/user_configs/config_options.pyi +2 -2
  251. metaflow-stubs/user_configs/config_parameters.pyi +4 -4
  252. metaflow-stubs/user_decorators/__init__.pyi +1 -1
  253. metaflow-stubs/user_decorators/common.pyi +1 -1
  254. metaflow-stubs/user_decorators/mutable_flow.pyi +3 -3
  255. metaflow-stubs/user_decorators/mutable_step.pyi +4 -4
  256. metaflow-stubs/user_decorators/user_flow_decorator.pyi +4 -4
  257. metaflow-stubs/user_decorators/user_step_decorator.pyi +4 -4
  258. {ob_metaflow_stubs-6.0.7.4.dist-info → ob_metaflow_stubs-6.0.8.1.dist-info}/METADATA +1 -1
  259. ob_metaflow_stubs-6.0.8.1.dist-info/RECORD +262 -0
  260. ob_metaflow_stubs-6.0.7.4.dist-info/RECORD +0 -262
  261. {ob_metaflow_stubs-6.0.7.4.dist-info → ob_metaflow_stubs-6.0.8.1.dist-info}/WHEEL +0 -0
  262. {ob_metaflow_stubs-6.0.7.4.dist-info → ob_metaflow_stubs-6.0.8.1.dist-info}/top_level.txt +0 -0
@@ -1,7 +1,7 @@
1
1
  ######################################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
3
  # MF version: 2.17.1.0+obcheckpoint(0.2.4);ob(v1) #
4
- # Generated on 2025-08-21T17:03:57.569240 #
4
+ # Generated on 2025-08-21T23:31:59.916830 #
5
5
  ######################################################################################################
6
6
 
7
7
  from __future__ import annotations
@@ -39,17 +39,17 @@ from .user_decorators.user_step_decorator import UserStepDecorator as UserStepDe
39
39
  from .user_decorators.user_step_decorator import StepMutator as StepMutator
40
40
  from .user_decorators.user_step_decorator import user_step_decorator as user_step_decorator
41
41
  from .user_decorators.user_flow_decorator import FlowMutator as FlowMutator
42
- from . import events as events
43
- from . import metaflow_git as metaflow_git
44
42
  from . import cards as cards
45
43
  from . import tuple_util as tuple_util
44
+ from . import metaflow_git as metaflow_git
45
+ from . import events as events
46
46
  from . import runner as runner
47
47
  from . import plugins as plugins
48
48
  from .mf_extensions.outerbounds.toplevel.global_aliases_for_metaflow_package import S3 as S3
49
49
  from . import includefile as includefile
50
50
  from .includefile import IncludeFile as IncludeFile
51
- from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
52
51
  from .plugins.pypi.parsers import conda_environment_yml_parser as conda_environment_yml_parser
52
+ from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
53
53
  from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
54
54
  from . import client as client
55
55
  from .client.core import namespace as namespace
@@ -167,160 +167,101 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
167
167
  """
168
168
  ...
169
169
 
170
- def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: typing.Optional[str] = None, image_pull_policy: str = 'KUBERNETES_IMAGE_PULL_POLICY', image_pull_secrets: typing.List[str] = [], service_account: str = 'METAFLOW_KUBERNETES_SERVICE_ACCOUNT', secrets: typing.Optional[typing.List[str]] = None, node_selector: typing.Union[typing.Dict[str, str], str, None] = None, namespace: str = 'METAFLOW_KUBERNETES_NAMESPACE', gpu: typing.Optional[int] = None, gpu_vendor: str = 'KUBERNETES_GPU_VENDOR', tolerations: typing.List[typing.Dict[str, str]] = [], labels: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_LABELS', annotations: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_ANNOTATIONS', use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = '/metaflow_temp', persistent_volume_claims: typing.Optional[typing.Dict[str, str]] = None, shared_memory: typing.Optional[int] = None, port: typing.Optional[int] = None, compute_pool: typing.Optional[str] = None, hostname_resolution_timeout: int = 600, qos: str = 'Burstable', security_context: typing.Optional[typing.Dict[str, typing.Any]] = 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]]]:
170
+ def ollama(*, models: list, backend: str, force_pull: bool, cache_update_policy: str, force_cache_update: bool, debug: bool, circuit_breaker_config: dict, timeout_config: dict) -> 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]]]:
171
171
  """
172
- Specifies that this step should execute on Kubernetes.
172
+ This decorator is used to run Ollama APIs as Metaflow task sidecars.
173
+
174
+ User code call
175
+ --------------
176
+ @ollama(
177
+ models=[...],
178
+ ...
179
+ )
180
+
181
+ Valid backend options
182
+ ---------------------
183
+ - 'local': Run as a separate process on the local task machine.
184
+ - (TODO) 'managed': Outerbounds hosts and selects compute provider.
185
+ - (TODO) 'remote': Spin up separate instance to serve Ollama models.
186
+
187
+ Valid model options
188
+ -------------------
189
+ Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
173
190
 
174
191
 
175
192
  Parameters
176
193
  ----------
177
- cpu : int, default 1
178
- Number of CPUs required for this step. If `@resources` is
179
- also present, the maximum value from all decorators is used.
180
- memory : int, default 4096
181
- Memory size (in MB) required for this step. If
182
- `@resources` is also present, the maximum value from all decorators is
183
- used.
184
- disk : int, default 10240
185
- Disk size (in MB) required for this step. If
186
- `@resources` is also present, the maximum value from all decorators is
187
- used.
188
- image : str, optional, default None
189
- Docker image to use when launching on Kubernetes. If not specified, and
190
- METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
191
- not, a default Docker image mapping to the current version of Python is used.
192
- image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
193
- If given, the imagePullPolicy to be applied to the Docker image of the step.
194
- image_pull_secrets: List[str], default []
195
- The default is extracted from METAFLOW_KUBERNETES_IMAGE_PULL_SECRETS.
196
- Kubernetes image pull secrets to use when pulling container images
197
- in Kubernetes.
198
- service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
199
- Kubernetes service account to use when launching pod in Kubernetes.
200
- secrets : List[str], optional, default None
201
- Kubernetes secrets to use when launching pod in Kubernetes. These
202
- secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
203
- in Metaflow configuration.
204
- node_selector: Union[Dict[str,str], str], optional, default None
205
- Kubernetes node selector(s) to apply to the pod running the task.
206
- Can be passed in as a comma separated string of values e.g.
207
- 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
208
- {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
209
- namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
210
- Kubernetes namespace to use when launching pod in Kubernetes.
211
- gpu : int, optional, default None
212
- Number of GPUs required for this step. A value of zero implies that
213
- the scheduled node should not have GPUs.
214
- gpu_vendor : str, default KUBERNETES_GPU_VENDOR
215
- The vendor of the GPUs to be used for this step.
216
- tolerations : List[Dict[str,str]], default []
217
- The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
218
- Kubernetes tolerations to use when launching pod in Kubernetes.
219
- labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
220
- Kubernetes labels to use when launching pod in Kubernetes.
221
- annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
222
- Kubernetes annotations to use when launching pod in Kubernetes.
223
- use_tmpfs : bool, default False
224
- This enables an explicit tmpfs mount for this step.
225
- tmpfs_tempdir : bool, default True
226
- sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
227
- tmpfs_size : int, optional, default: None
228
- The value for the size (in MiB) of the tmpfs mount for this step.
229
- This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
230
- memory allocated for this step.
231
- tmpfs_path : str, optional, default /metaflow_temp
232
- Path to tmpfs mount for this step.
233
- persistent_volume_claims : Dict[str, str], optional, default None
234
- A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
235
- volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
236
- shared_memory: int, optional
237
- Shared memory size (in MiB) required for this step
238
- port: int, optional
239
- Port number to specify in the Kubernetes job object
240
- compute_pool : str, optional, default None
241
- Compute pool to be used for for this step.
242
- If not specified, any accessible compute pool within the perimeter is used.
243
- hostname_resolution_timeout: int, default 10 * 60
244
- Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
245
- Only applicable when @parallel is used.
246
- qos: str, default: Burstable
247
- Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
248
-
249
- security_context: Dict[str, Any], optional, default None
250
- Container security context. Applies to the task container. Allows the following keys:
251
- - privileged: bool, optional, default None
252
- - allow_privilege_escalation: bool, optional, default None
253
- - run_as_user: int, optional, default None
254
- - run_as_group: int, optional, default None
255
- - run_as_non_root: bool, optional, default None
256
- """
257
- ...
258
-
259
- @typing.overload
260
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
261
- """
262
- Internal decorator to support Fast bakery
263
- """
264
- ...
265
-
266
- @typing.overload
267
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
268
- ...
269
-
270
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
271
- """
272
- Internal decorator to support Fast bakery
194
+ models: list[str]
195
+ List of Ollama containers running models in sidecars.
196
+ backend: str
197
+ Determines where and how to run the Ollama process.
198
+ force_pull: bool
199
+ Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
200
+ cache_update_policy: str
201
+ Cache update policy: "auto", "force", or "never".
202
+ force_cache_update: bool
203
+ Simple override for "force" cache update policy.
204
+ debug: bool
205
+ Whether to turn on verbose debugging logs.
206
+ circuit_breaker_config: dict
207
+ Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
208
+ timeout_config: dict
209
+ Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
273
210
  """
274
211
  ...
275
212
 
276
213
  @typing.overload
277
- 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]]]:
214
+ 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]]]:
278
215
  """
279
- Specifies the PyPI packages for the step.
216
+ Specifies the number of times the task corresponding
217
+ to a step needs to be retried.
280
218
 
281
- Information in this decorator will augment any
282
- attributes set in the `@pyi_base` flow-level decorator. Hence,
283
- you can use `@pypi_base` to set packages required by all
284
- steps and use `@pypi` to specify step-specific overrides.
219
+ This decorator is useful for handling transient errors, such as networking issues.
220
+ If your task contains operations that can't be retried safely, e.g. database updates,
221
+ it is advisable to annotate it with `@retry(times=0)`.
222
+
223
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
224
+ decorator will execute a no-op task after all retries have been exhausted,
225
+ ensuring that the flow execution can continue.
285
226
 
286
227
 
287
228
  Parameters
288
229
  ----------
289
- packages : Dict[str, str], default: {}
290
- Packages to use for this step. The key is the name of the package
291
- and the value is the version to use.
292
- python : str, optional, default: None
293
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
294
- that the version used will correspond to the version of the Python interpreter used to start the run.
230
+ times : int, default 3
231
+ Number of times to retry this task.
232
+ minutes_between_retries : int, default 2
233
+ Number of minutes between retries.
295
234
  """
296
235
  ...
297
236
 
298
237
  @typing.overload
299
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
238
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
300
239
  ...
301
240
 
302
241
  @typing.overload
303
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
242
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
304
243
  ...
305
244
 
306
- 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):
245
+ 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):
307
246
  """
308
- Specifies the PyPI packages for the step.
247
+ Specifies the number of times the task corresponding
248
+ to a step needs to be retried.
309
249
 
310
- Information in this decorator will augment any
311
- attributes set in the `@pyi_base` flow-level decorator. Hence,
312
- you can use `@pypi_base` to set packages required by all
313
- steps and use `@pypi` to specify step-specific overrides.
250
+ This decorator is useful for handling transient errors, such as networking issues.
251
+ If your task contains operations that can't be retried safely, e.g. database updates,
252
+ it is advisable to annotate it with `@retry(times=0)`.
253
+
254
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
255
+ decorator will execute a no-op task after all retries have been exhausted,
256
+ ensuring that the flow execution can continue.
314
257
 
315
258
 
316
259
  Parameters
317
260
  ----------
318
- packages : Dict[str, str], default: {}
319
- Packages to use for this step. The key is the name of the package
320
- and the value is the version to use.
321
- python : str, optional, default: None
322
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
323
- that the version used will correspond to the version of the Python interpreter used to start the run.
261
+ times : int, default 3
262
+ Number of times to retry this task.
263
+ minutes_between_retries : int, default 2
264
+ Number of minutes between retries.
324
265
  """
325
266
  ...
326
267
 
@@ -375,257 +316,51 @@ def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
375
316
  """
376
317
  ...
377
318
 
378
- def vllm(*, model: str, backend: str, openai_api_server: bool, debug: bool, card_refresh_interval: int, max_retries: int, retry_alert_frequency: int, engine_args: dict) -> 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]]]:
319
+ @typing.overload
320
+ 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]]]:
379
321
  """
380
- This decorator is used to run vllm APIs as Metaflow task sidecars.
381
-
382
- User code call
383
- --------------
384
- @vllm(
385
- model="...",
386
- ...
387
- )
322
+ Enables loading / saving of models within a step.
388
323
 
389
- Valid backend options
390
- ---------------------
391
- - 'local': Run as a separate process on the local task machine.
324
+ > Examples
325
+ - Saving Models
326
+ ```python
327
+ @model
328
+ @step
329
+ def train(self):
330
+ # current.model.save returns a dictionary reference to the model saved
331
+ self.my_model = current.model.save(
332
+ path_to_my_model,
333
+ label="my_model",
334
+ metadata={
335
+ "epochs": 10,
336
+ "batch-size": 32,
337
+ "learning-rate": 0.001,
338
+ }
339
+ )
340
+ self.next(self.test)
392
341
 
393
- Valid model options
394
- -------------------
395
- Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
342
+ @model(load="my_model")
343
+ @step
344
+ def test(self):
345
+ # `current.model.loaded` returns a dictionary of the loaded models
346
+ # where the key is the name of the artifact and the value is the path to the model
347
+ print(os.listdir(current.model.loaded["my_model"]))
348
+ self.next(self.end)
349
+ ```
396
350
 
397
- NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
398
- If you need multiple models, you must create multiple @vllm decorators.
399
-
400
-
401
- Parameters
402
- ----------
403
- model: str
404
- HuggingFace model identifier to be served by vLLM.
405
- backend: str
406
- Determines where and how to run the vLLM process.
407
- openai_api_server: bool
408
- Whether to use OpenAI-compatible API server mode (subprocess) instead of native engine.
409
- Default is False (uses native engine).
410
- Set to True for backward compatibility with existing code.
411
- debug: bool
412
- Whether to turn on verbose debugging logs.
413
- card_refresh_interval: int
414
- Interval in seconds for refreshing the vLLM status card.
415
- Only used when openai_api_server=True.
416
- max_retries: int
417
- Maximum number of retries checking for vLLM server startup.
418
- Only used when openai_api_server=True.
419
- retry_alert_frequency: int
420
- Frequency of alert logs for vLLM server startup retries.
421
- Only used when openai_api_server=True.
422
- engine_args : dict
423
- Additional keyword arguments to pass to the vLLM engine.
424
- For example, `tensor_parallel_size=2`.
425
- """
426
- ...
427
-
428
- @typing.overload
429
- 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]]]:
430
- """
431
- Specifies the number of times the task corresponding
432
- to a step needs to be retried.
433
-
434
- This decorator is useful for handling transient errors, such as networking issues.
435
- If your task contains operations that can't be retried safely, e.g. database updates,
436
- it is advisable to annotate it with `@retry(times=0)`.
437
-
438
- This can be used in conjunction with the `@catch` decorator. The `@catch`
439
- decorator will execute a no-op task after all retries have been exhausted,
440
- ensuring that the flow execution can continue.
441
-
442
-
443
- Parameters
444
- ----------
445
- times : int, default 3
446
- Number of times to retry this task.
447
- minutes_between_retries : int, default 2
448
- Number of minutes between retries.
449
- """
450
- ...
451
-
452
- @typing.overload
453
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
454
- ...
455
-
456
- @typing.overload
457
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
458
- ...
459
-
460
- 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):
461
- """
462
- Specifies the number of times the task corresponding
463
- to a step needs to be retried.
464
-
465
- This decorator is useful for handling transient errors, such as networking issues.
466
- If your task contains operations that can't be retried safely, e.g. database updates,
467
- it is advisable to annotate it with `@retry(times=0)`.
468
-
469
- This can be used in conjunction with the `@catch` decorator. The `@catch`
470
- decorator will execute a no-op task after all retries have been exhausted,
471
- ensuring that the flow execution can continue.
472
-
473
-
474
- Parameters
475
- ----------
476
- times : int, default 3
477
- Number of times to retry this task.
478
- minutes_between_retries : int, default 2
479
- Number of minutes between retries.
480
- """
481
- ...
482
-
483
- @typing.overload
484
- 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]]]:
485
- """
486
- Specifies environment variables to be set prior to the execution of a step.
487
-
488
-
489
- Parameters
490
- ----------
491
- vars : Dict[str, str], default {}
492
- Dictionary of environment variables to set.
493
- """
494
- ...
495
-
496
- @typing.overload
497
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
498
- ...
499
-
500
- @typing.overload
501
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
502
- ...
503
-
504
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
505
- """
506
- Specifies environment variables to be set prior to the execution of a step.
507
-
508
-
509
- Parameters
510
- ----------
511
- vars : Dict[str, str], default {}
512
- Dictionary of environment variables to set.
513
- """
514
- ...
515
-
516
- @typing.overload
517
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
518
- """
519
- Decorator prototype for all step decorators. This function gets specialized
520
- and imported for all decorators types by _import_plugin_decorators().
521
- """
522
- ...
523
-
524
- @typing.overload
525
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
526
- ...
527
-
528
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
529
- """
530
- Decorator prototype for all step decorators. This function gets specialized
531
- and imported for all decorators types by _import_plugin_decorators().
532
- """
533
- ...
534
-
535
- @typing.overload
536
- 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]]]:
537
- """
538
- Creates a human-readable report, a Metaflow Card, after this step completes.
539
-
540
- Note that you may add multiple `@card` decorators in a step with different parameters.
541
-
542
-
543
- Parameters
544
- ----------
545
- type : str, default 'default'
546
- Card type.
547
- id : str, optional, default None
548
- If multiple cards are present, use this id to identify this card.
549
- options : Dict[str, Any], default {}
550
- Options passed to the card. The contents depend on the card type.
551
- timeout : int, default 45
552
- Interrupt reporting if it takes more than this many seconds.
553
- """
554
- ...
555
-
556
- @typing.overload
557
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
558
- ...
559
-
560
- @typing.overload
561
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
562
- ...
563
-
564
- 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):
565
- """
566
- Creates a human-readable report, a Metaflow Card, after this step completes.
567
-
568
- Note that you may add multiple `@card` decorators in a step with different parameters.
569
-
570
-
571
- Parameters
572
- ----------
573
- type : str, default 'default'
574
- Card type.
575
- id : str, optional, default None
576
- If multiple cards are present, use this id to identify this card.
577
- options : Dict[str, Any], default {}
578
- Options passed to the card. The contents depend on the card type.
579
- timeout : int, default 45
580
- Interrupt reporting if it takes more than this many seconds.
581
- """
582
- ...
583
-
584
- @typing.overload
585
- 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]]]:
586
- """
587
- Enables loading / saving of models within a step.
588
-
589
- > Examples
590
- - Saving Models
591
- ```python
592
- @model
593
- @step
594
- def train(self):
595
- # current.model.save returns a dictionary reference to the model saved
596
- self.my_model = current.model.save(
597
- path_to_my_model,
598
- label="my_model",
599
- metadata={
600
- "epochs": 10,
601
- "batch-size": 32,
602
- "learning-rate": 0.001,
603
- }
604
- )
605
- self.next(self.test)
606
-
607
- @model(load="my_model")
608
- @step
609
- def test(self):
610
- # `current.model.loaded` returns a dictionary of the loaded models
611
- # where the key is the name of the artifact and the value is the path to the model
612
- print(os.listdir(current.model.loaded["my_model"]))
613
- self.next(self.end)
614
- ```
615
-
616
- - Loading models
617
- ```python
618
- @step
619
- def train(self):
620
- # current.model.load returns the path to the model loaded
621
- checkpoint_path = current.model.load(
622
- self.checkpoint_key,
623
- )
624
- model_path = current.model.load(
625
- self.model,
626
- )
627
- self.next(self.test)
628
- ```
351
+ - Loading models
352
+ ```python
353
+ @step
354
+ def train(self):
355
+ # current.model.load returns the path to the model loaded
356
+ checkpoint_path = current.model.load(
357
+ self.checkpoint_key,
358
+ )
359
+ model_path = current.model.load(
360
+ self.model,
361
+ )
362
+ self.next(self.test)
363
+ ```
629
364
 
630
365
 
631
366
  Parameters
@@ -710,167 +445,53 @@ def model(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
710
445
  """
711
446
  ...
712
447
 
713
- @typing.overload
714
- def app_deploy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
715
- """
716
- Decorator prototype for all step decorators. This function gets specialized
717
- and imported for all decorators types by _import_plugin_decorators().
718
- """
719
- ...
720
-
721
- @typing.overload
722
- def app_deploy(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
723
- ...
724
-
725
- def app_deploy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
726
- """
727
- Decorator prototype for all step decorators. This function gets specialized
728
- and imported for all decorators types by _import_plugin_decorators().
729
- """
730
- ...
731
-
732
- def s3_proxy(*, integration_name: typing.Optional[str] = None, write_mode: typing.Optional[str] = None, debug: typing.Optional[bool] = 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]]]:
448
+ def vllm(*, model: str, backend: str, openai_api_server: bool, debug: bool, card_refresh_interval: int, max_retries: int, retry_alert_frequency: int, engine_args: dict) -> 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]]]:
733
449
  """
734
- S3 Proxy decorator for routing S3 requests through a local proxy service.
450
+ This decorator is used to run vllm APIs as Metaflow task sidecars.
735
451
 
736
-
737
- Parameters
738
- ----------
739
- integration_name : str, optional
740
- Name of the S3 proxy integration. If not specified, will use the only
741
- available S3 proxy integration in the namespace (fails if multiple exist).
742
- write_mode : str, optional
743
- The desired behavior during write operations to target (origin) S3 bucket.
744
- allowed options are:
745
- "origin-and-cache" -> write to both the target S3 bucket and local object
746
- storage
747
- "origin" -> only write to the target S3 bucket
748
- "cache" -> only write to the object storage service used for caching
749
- debug : bool, optional
750
- Enable debug logging for proxy operations.
751
- """
752
- ...
753
-
754
- @typing.overload
755
- 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]]]:
756
- """
757
- Specifies the Conda environment for the step.
758
-
759
- Information in this decorator will augment any
760
- attributes set in the `@conda_base` flow-level decorator. Hence,
761
- you can use `@conda_base` to set packages required by all
762
- steps and use `@conda` to specify step-specific overrides.
763
-
764
-
765
- Parameters
766
- ----------
767
- packages : Dict[str, str], default {}
768
- Packages to use for this step. The key is the name of the package
769
- and the value is the version to use.
770
- libraries : Dict[str, str], default {}
771
- Supported for backward compatibility. When used with packages, packages will take precedence.
772
- python : str, optional, default None
773
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
774
- that the version used will correspond to the version of the Python interpreter used to start the run.
775
- disabled : bool, default False
776
- If set to True, disables @conda.
777
- """
778
- ...
779
-
780
- @typing.overload
781
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
782
- ...
783
-
784
- @typing.overload
785
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
786
- ...
787
-
788
- 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):
789
- """
790
- Specifies the Conda environment for the step.
791
-
792
- Information in this decorator will augment any
793
- attributes set in the `@conda_base` flow-level decorator. Hence,
794
- you can use `@conda_base` to set packages required by all
795
- steps and use `@conda` to specify step-specific overrides.
796
-
797
-
798
- Parameters
799
- ----------
800
- packages : Dict[str, str], default {}
801
- Packages to use for this step. The key is the name of the package
802
- and the value is the version to use.
803
- libraries : Dict[str, str], default {}
804
- Supported for backward compatibility. When used with packages, packages will take precedence.
805
- python : str, optional, default None
806
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
807
- that the version used will correspond to the version of the Python interpreter used to start the run.
808
- disabled : bool, default False
809
- If set to True, disables @conda.
810
- """
811
- ...
812
-
813
- def ollama(*, models: list, backend: str, force_pull: bool, cache_update_policy: str, force_cache_update: bool, debug: bool, circuit_breaker_config: dict, timeout_config: dict) -> 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]]]:
814
- """
815
- This decorator is used to run Ollama APIs as Metaflow task sidecars.
816
-
817
- User code call
818
- --------------
819
- @ollama(
820
- models=[...],
821
- ...
822
- )
452
+ User code call
453
+ --------------
454
+ @vllm(
455
+ model="...",
456
+ ...
457
+ )
823
458
 
824
459
  Valid backend options
825
460
  ---------------------
826
461
  - 'local': Run as a separate process on the local task machine.
827
- - (TODO) 'managed': Outerbounds hosts and selects compute provider.
828
- - (TODO) 'remote': Spin up separate instance to serve Ollama models.
829
462
 
830
463
  Valid model options
831
464
  -------------------
832
- Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
465
+ Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
466
+
467
+ NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
468
+ If you need multiple models, you must create multiple @vllm decorators.
833
469
 
834
470
 
835
471
  Parameters
836
472
  ----------
837
- models: list[str]
838
- List of Ollama containers running models in sidecars.
473
+ model: str
474
+ HuggingFace model identifier to be served by vLLM.
839
475
  backend: str
840
- Determines where and how to run the Ollama process.
841
- force_pull: bool
842
- Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
843
- cache_update_policy: str
844
- Cache update policy: "auto", "force", or "never".
845
- force_cache_update: bool
846
- Simple override for "force" cache update policy.
476
+ Determines where and how to run the vLLM process.
477
+ openai_api_server: bool
478
+ Whether to use OpenAI-compatible API server mode (subprocess) instead of native engine.
479
+ Default is False (uses native engine).
480
+ Set to True for backward compatibility with existing code.
847
481
  debug: bool
848
482
  Whether to turn on verbose debugging logs.
849
- circuit_breaker_config: dict
850
- Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
851
- timeout_config: dict
852
- Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
853
- """
854
- ...
855
-
856
- @typing.overload
857
- def nebius_s3_proxy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
858
- """
859
- Nebius-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
860
- It exists to make it easier for users to know that this decorator should only be used with
861
- a Neo Cloud like Nebius.
862
- """
863
- ...
864
-
865
- @typing.overload
866
- def nebius_s3_proxy(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
867
- ...
868
-
869
- def nebius_s3_proxy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
870
- """
871
- Nebius-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
872
- It exists to make it easier for users to know that this decorator should only be used with
873
- a Neo Cloud like Nebius.
483
+ card_refresh_interval: int
484
+ Interval in seconds for refreshing the vLLM status card.
485
+ Only used when openai_api_server=True.
486
+ max_retries: int
487
+ Maximum number of retries checking for vLLM server startup.
488
+ Only used when openai_api_server=True.
489
+ retry_alert_frequency: int
490
+ Frequency of alert logs for vLLM server startup retries.
491
+ Only used when openai_api_server=True.
492
+ engine_args : dict
493
+ Additional keyword arguments to pass to the vLLM engine.
494
+ For example, `tensor_parallel_size=2`.
874
495
  """
875
496
  ...
876
497
 
@@ -933,106 +554,6 @@ def timeout(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None],
933
554
  """
934
555
  ...
935
556
 
936
- @typing.overload
937
- def coreweave_s3_proxy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
938
- """
939
- CoreWeave-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
940
- It exists to make it easier for users to know that this decorator should only be used with
941
- a Neo Cloud like CoreWeave.
942
- """
943
- ...
944
-
945
- @typing.overload
946
- def coreweave_s3_proxy(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
947
- ...
948
-
949
- def coreweave_s3_proxy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
950
- """
951
- CoreWeave-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
952
- It exists to make it easier for users to know that this decorator should only be used with
953
- a Neo Cloud like CoreWeave.
954
- """
955
- ...
956
-
957
- @typing.overload
958
- 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]]]:
959
- """
960
- Specifies the resources needed when executing this step.
961
-
962
- Use `@resources` to specify the resource requirements
963
- independently of the specific compute layer (`@batch`, `@kubernetes`).
964
-
965
- You can choose the compute layer on the command line by executing e.g.
966
- ```
967
- python myflow.py run --with batch
968
- ```
969
- or
970
- ```
971
- python myflow.py run --with kubernetes
972
- ```
973
- which executes the flow on the desired system using the
974
- requirements specified in `@resources`.
975
-
976
-
977
- Parameters
978
- ----------
979
- cpu : int, default 1
980
- Number of CPUs required for this step.
981
- gpu : int, optional, default None
982
- Number of GPUs required for this step.
983
- disk : int, optional, default None
984
- Disk size (in MB) required for this step. Only applies on Kubernetes.
985
- memory : int, default 4096
986
- Memory size (in MB) required for this step.
987
- shared_memory : int, optional, default None
988
- The value for the size (in MiB) of the /dev/shm volume for this step.
989
- This parameter maps to the `--shm-size` option in Docker.
990
- """
991
- ...
992
-
993
- @typing.overload
994
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
995
- ...
996
-
997
- @typing.overload
998
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
999
- ...
1000
-
1001
- 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):
1002
- """
1003
- Specifies the resources needed when executing this step.
1004
-
1005
- Use `@resources` to specify the resource requirements
1006
- independently of the specific compute layer (`@batch`, `@kubernetes`).
1007
-
1008
- You can choose the compute layer on the command line by executing e.g.
1009
- ```
1010
- python myflow.py run --with batch
1011
- ```
1012
- or
1013
- ```
1014
- python myflow.py run --with kubernetes
1015
- ```
1016
- which executes the flow on the desired system using the
1017
- requirements specified in `@resources`.
1018
-
1019
-
1020
- Parameters
1021
- ----------
1022
- cpu : int, default 1
1023
- Number of CPUs required for this step.
1024
- gpu : int, optional, default None
1025
- Number of GPUs required for this step.
1026
- disk : int, optional, default None
1027
- Disk size (in MB) required for this step. Only applies on Kubernetes.
1028
- memory : int, default 4096
1029
- Memory size (in MB) required for this step.
1030
- shared_memory : int, optional, default None
1031
- The value for the size (in MiB) of the /dev/shm volume for this step.
1032
- This parameter maps to the `--shm-size` option in Docker.
1033
- """
1034
- ...
1035
-
1036
557
  @typing.overload
1037
558
  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]]]:
1038
559
  """
@@ -1180,85 +701,11 @@ def checkpoint(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None
1180
701
  """
1181
702
  ...
1182
703
 
1183
- @typing.overload
1184
- def secrets(*, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], role: 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]]]:
704
+ 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]]]:
1185
705
  """
1186
- Specifies secrets to be retrieved and injected as environment variables prior to
1187
- the execution of a step.
706
+ Decorator that helps cache, version and store models/datasets from huggingface hub.
1188
707
 
1189
-
1190
- Parameters
1191
- ----------
1192
- sources : List[Union[str, Dict[str, Any]]], default: []
1193
- List of secret specs, defining how the secrets are to be retrieved
1194
- role : str, optional, default: None
1195
- Role to use for fetching secrets
1196
- """
1197
- ...
1198
-
1199
- @typing.overload
1200
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1201
- ...
1202
-
1203
- @typing.overload
1204
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1205
- ...
1206
-
1207
- 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]]] = [], role: typing.Optional[str] = None):
1208
- """
1209
- Specifies secrets to be retrieved and injected as environment variables prior to
1210
- the execution of a step.
1211
-
1212
-
1213
- Parameters
1214
- ----------
1215
- sources : List[Union[str, Dict[str, Any]]], default: []
1216
- List of secret specs, defining how the secrets are to be retrieved
1217
- role : str, optional, default: None
1218
- Role to use for fetching secrets
1219
- """
1220
- ...
1221
-
1222
- @typing.overload
1223
- def test_append_card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1224
- """
1225
- A simple decorator that demonstrates using CardDecoratorInjector
1226
- to inject a card and render simple markdown content.
1227
- """
1228
- ...
1229
-
1230
- @typing.overload
1231
- def test_append_card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1232
- ...
1233
-
1234
- def test_append_card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1235
- """
1236
- A simple decorator that demonstrates using CardDecoratorInjector
1237
- to inject a card and render simple markdown content.
1238
- """
1239
- ...
1240
-
1241
- 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]]]:
1242
- """
1243
- Specifies that this step should execute on DGX cloud.
1244
-
1245
-
1246
- Parameters
1247
- ----------
1248
- gpu : int
1249
- Number of GPUs to use.
1250
- gpu_type : str
1251
- Type of Nvidia GPU to use.
1252
- queue_timeout : int
1253
- Time to keep the job in NVCF's queue.
1254
- """
1255
- ...
1256
-
1257
- 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]]]:
1258
- """
1259
- Decorator that helps cache, version and store models/datasets from huggingface hub.
1260
-
1261
- > Examples
708
+ > Examples
1262
709
 
1263
710
  **Usage: creating references of models from huggingface that may be loaded in downstream steps**
1264
711
  ```python
@@ -1334,126 +781,143 @@ def huggingface_hub(*, temp_dir_root: typing.Optional[str] = None, load: typing.
1334
781
  """
1335
782
  ...
1336
783
 
1337
- def nvct(*, gpu: int, gpu_type: 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]]]:
784
+ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: typing.Optional[str] = None, image_pull_policy: str = 'KUBERNETES_IMAGE_PULL_POLICY', image_pull_secrets: typing.List[str] = [], service_account: str = 'METAFLOW_KUBERNETES_SERVICE_ACCOUNT', secrets: typing.Optional[typing.List[str]] = None, node_selector: typing.Union[typing.Dict[str, str], str, None] = None, namespace: str = 'METAFLOW_KUBERNETES_NAMESPACE', gpu: typing.Optional[int] = None, gpu_vendor: str = 'KUBERNETES_GPU_VENDOR', tolerations: typing.List[typing.Dict[str, str]] = [], labels: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_LABELS', annotations: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_ANNOTATIONS', use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = '/metaflow_temp', persistent_volume_claims: typing.Optional[typing.Dict[str, str]] = None, shared_memory: typing.Optional[int] = None, port: typing.Optional[int] = None, compute_pool: typing.Optional[str] = None, hostname_resolution_timeout: int = 600, qos: str = 'Burstable', security_context: typing.Optional[typing.Dict[str, typing.Any]] = 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]]]:
1338
785
  """
1339
- Specifies that this step should execute on DGX cloud.
786
+ Specifies that this step should execute on Kubernetes.
1340
787
 
1341
788
 
1342
789
  Parameters
1343
790
  ----------
1344
- gpu : int
1345
- Number of GPUs to use.
1346
- gpu_type : str
1347
- Type of Nvidia GPU to use.
791
+ cpu : int, default 1
792
+ Number of CPUs required for this step. If `@resources` is
793
+ also present, the maximum value from all decorators is used.
794
+ memory : int, default 4096
795
+ Memory size (in MB) required for this step. If
796
+ `@resources` is also present, the maximum value from all decorators is
797
+ used.
798
+ disk : int, default 10240
799
+ Disk size (in MB) required for this step. If
800
+ `@resources` is also present, the maximum value from all decorators is
801
+ used.
802
+ image : str, optional, default None
803
+ Docker image to use when launching on Kubernetes. If not specified, and
804
+ METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
805
+ not, a default Docker image mapping to the current version of Python is used.
806
+ image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
807
+ If given, the imagePullPolicy to be applied to the Docker image of the step.
808
+ image_pull_secrets: List[str], default []
809
+ The default is extracted from METAFLOW_KUBERNETES_IMAGE_PULL_SECRETS.
810
+ Kubernetes image pull secrets to use when pulling container images
811
+ in Kubernetes.
812
+ service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
813
+ Kubernetes service account to use when launching pod in Kubernetes.
814
+ secrets : List[str], optional, default None
815
+ Kubernetes secrets to use when launching pod in Kubernetes. These
816
+ secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
817
+ in Metaflow configuration.
818
+ node_selector: Union[Dict[str,str], str], optional, default None
819
+ Kubernetes node selector(s) to apply to the pod running the task.
820
+ Can be passed in as a comma separated string of values e.g.
821
+ 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
822
+ {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
823
+ namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
824
+ Kubernetes namespace to use when launching pod in Kubernetes.
825
+ gpu : int, optional, default None
826
+ Number of GPUs required for this step. A value of zero implies that
827
+ the scheduled node should not have GPUs.
828
+ gpu_vendor : str, default KUBERNETES_GPU_VENDOR
829
+ The vendor of the GPUs to be used for this step.
830
+ tolerations : List[Dict[str,str]], default []
831
+ The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
832
+ Kubernetes tolerations to use when launching pod in Kubernetes.
833
+ labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
834
+ Kubernetes labels to use when launching pod in Kubernetes.
835
+ annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
836
+ Kubernetes annotations to use when launching pod in Kubernetes.
837
+ use_tmpfs : bool, default False
838
+ This enables an explicit tmpfs mount for this step.
839
+ tmpfs_tempdir : bool, default True
840
+ sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
841
+ tmpfs_size : int, optional, default: None
842
+ The value for the size (in MiB) of the tmpfs mount for this step.
843
+ This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
844
+ memory allocated for this step.
845
+ tmpfs_path : str, optional, default /metaflow_temp
846
+ Path to tmpfs mount for this step.
847
+ persistent_volume_claims : Dict[str, str], optional, default None
848
+ A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
849
+ volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
850
+ shared_memory: int, optional
851
+ Shared memory size (in MiB) required for this step
852
+ port: int, optional
853
+ Port number to specify in the Kubernetes job object
854
+ compute_pool : str, optional, default None
855
+ Compute pool to be used for for this step.
856
+ If not specified, any accessible compute pool within the perimeter is used.
857
+ hostname_resolution_timeout: int, default 10 * 60
858
+ Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
859
+ Only applicable when @parallel is used.
860
+ qos: str, default: Burstable
861
+ Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
862
+
863
+ security_context: Dict[str, Any], optional, default None
864
+ Container security context. Applies to the task container. Allows the following keys:
865
+ - privileged: bool, optional, default None
866
+ - allow_privilege_escalation: bool, optional, default None
867
+ - run_as_user: int, optional, default None
868
+ - run_as_group: int, optional, default None
869
+ - run_as_non_root: bool, optional, default None
1348
870
  """
1349
871
  ...
1350
872
 
1351
- 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]]:
873
+ @typing.overload
874
+ 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]]]:
1352
875
  """
1353
- 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)
1354
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1355
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1356
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
1357
- starts only after all sensors finish.
876
+ Specifies environment variables to be set prior to the execution of a step.
1358
877
 
1359
878
 
1360
879
  Parameters
1361
880
  ----------
1362
- timeout : int
1363
- Time, in seconds before the task times out and fails. (Default: 3600)
1364
- poke_interval : int
1365
- Time in seconds that the job should wait in between each try. (Default: 60)
1366
- mode : str
1367
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1368
- exponential_backoff : bool
1369
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1370
- pool : str
1371
- the slot pool this task should run in,
1372
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1373
- soft_fail : bool
1374
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1375
- name : str
1376
- Name of the sensor on Airflow
1377
- description : str
1378
- Description of sensor in the Airflow UI
1379
- bucket_key : Union[str, List[str]]
1380
- The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1381
- When it's specified as a full s3:// url, please leave `bucket_name` as None
1382
- bucket_name : str
1383
- Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1384
- When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1385
- wildcard_match : bool
1386
- whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1387
- aws_conn_id : str
1388
- a reference to the s3 connection on Airflow. (Default: None)
1389
- verify : bool
1390
- Whether or not to verify SSL certificates for S3 connection. (Default: None)
881
+ vars : Dict[str, str], default {}
882
+ Dictionary of environment variables to set.
1391
883
  """
1392
884
  ...
1393
885
 
1394
886
  @typing.overload
1395
- def schedule(*, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1396
- """
1397
- Specifies the times when the flow should be run when running on a
1398
- production scheduler.
1399
-
1400
-
1401
- Parameters
1402
- ----------
1403
- hourly : bool, default False
1404
- Run the workflow hourly.
1405
- daily : bool, default True
1406
- Run the workflow daily.
1407
- weekly : bool, default False
1408
- Run the workflow weekly.
1409
- cron : str, optional, default None
1410
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1411
- specified by this expression.
1412
- timezone : str, optional, default None
1413
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1414
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1415
- """
887
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1416
888
  ...
1417
889
 
1418
890
  @typing.overload
1419
- def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
891
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1420
892
  ...
1421
893
 
1422
- def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None):
894
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
1423
895
  """
1424
- Specifies the times when the flow should be run when running on a
1425
- production scheduler.
896
+ Specifies environment variables to be set prior to the execution of a step.
1426
897
 
1427
898
 
1428
899
  Parameters
1429
900
  ----------
1430
- hourly : bool, default False
1431
- Run the workflow hourly.
1432
- daily : bool, default True
1433
- Run the workflow daily.
1434
- weekly : bool, default False
1435
- Run the workflow weekly.
1436
- cron : str, optional, default None
1437
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1438
- specified by this expression.
1439
- timezone : str, optional, default None
1440
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1441
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
901
+ vars : Dict[str, str], default {}
902
+ Dictionary of environment variables to set.
1442
903
  """
1443
904
  ...
1444
905
 
1445
906
  @typing.overload
1446
- def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
907
+ 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]]]:
1447
908
  """
1448
- Specifies the PyPI packages for all steps of the flow.
909
+ Specifies the PyPI packages for the step.
1449
910
 
1450
- Use `@pypi_base` to set common packages required by all
911
+ Information in this decorator will augment any
912
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
913
+ you can use `@pypi_base` to set packages required by all
1451
914
  steps and use `@pypi` to specify step-specific overrides.
1452
915
 
916
+
1453
917
  Parameters
1454
918
  ----------
1455
919
  packages : Dict[str, str], default: {}
1456
- Packages to use for this flow. The key is the name of the package
920
+ Packages to use for this step. The key is the name of the package
1457
921
  and the value is the version to use.
1458
922
  python : str, optional, default: None
1459
923
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
@@ -1462,107 +926,596 @@ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[s
1462
926
  ...
1463
927
 
1464
928
  @typing.overload
1465
- def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
929
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1466
930
  ...
1467
931
 
1468
- def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1469
- """
1470
- Specifies the PyPI packages for all steps of the flow.
932
+ @typing.overload
933
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
934
+ ...
935
+
936
+ 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):
937
+ """
938
+ Specifies the PyPI packages for the step.
939
+
940
+ Information in this decorator will augment any
941
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
942
+ you can use `@pypi_base` to set packages required by all
943
+ steps and use `@pypi` to specify step-specific overrides.
944
+
945
+
946
+ Parameters
947
+ ----------
948
+ packages : Dict[str, str], default: {}
949
+ Packages to use for this step. The key is the name of the package
950
+ and the value is the version to use.
951
+ python : str, optional, default: None
952
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
953
+ that the version used will correspond to the version of the Python interpreter used to start the run.
954
+ """
955
+ ...
956
+
957
+ 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]]]:
958
+ """
959
+ Specifies that this step should execute on DGX cloud.
960
+
961
+
962
+ Parameters
963
+ ----------
964
+ gpu : int
965
+ Number of GPUs to use.
966
+ gpu_type : str
967
+ Type of Nvidia GPU to use.
968
+ queue_timeout : int
969
+ Time to keep the job in NVCF's queue.
970
+ """
971
+ ...
972
+
973
+ @typing.overload
974
+ def test_append_card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
975
+ """
976
+ A simple decorator that demonstrates using CardDecoratorInjector
977
+ to inject a card and render simple markdown content.
978
+ """
979
+ ...
980
+
981
+ @typing.overload
982
+ def test_append_card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
983
+ ...
984
+
985
+ def test_append_card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
986
+ """
987
+ A simple decorator that demonstrates using CardDecoratorInjector
988
+ to inject a card and render simple markdown content.
989
+ """
990
+ ...
991
+
992
+ @typing.overload
993
+ def nebius_s3_proxy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
994
+ """
995
+ Nebius-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
996
+ It exists to make it easier for users to know that this decorator should only be used with
997
+ a Neo Cloud like Nebius.
998
+ """
999
+ ...
1000
+
1001
+ @typing.overload
1002
+ def nebius_s3_proxy(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1003
+ ...
1004
+
1005
+ def nebius_s3_proxy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1006
+ """
1007
+ Nebius-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
1008
+ It exists to make it easier for users to know that this decorator should only be used with
1009
+ a Neo Cloud like Nebius.
1010
+ """
1011
+ ...
1012
+
1013
+ def nvct(*, gpu: int, gpu_type: 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]]]:
1014
+ """
1015
+ Specifies that this step should execute on DGX cloud.
1016
+
1017
+
1018
+ Parameters
1019
+ ----------
1020
+ gpu : int
1021
+ Number of GPUs to use.
1022
+ gpu_type : str
1023
+ Type of Nvidia GPU to use.
1024
+ """
1025
+ ...
1026
+
1027
+ @typing.overload
1028
+ def secrets(*, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], role: 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]]]:
1029
+ """
1030
+ Specifies secrets to be retrieved and injected as environment variables prior to
1031
+ the execution of a step.
1032
+
1033
+
1034
+ Parameters
1035
+ ----------
1036
+ sources : List[Union[str, Dict[str, Any]]], default: []
1037
+ List of secret specs, defining how the secrets are to be retrieved
1038
+ role : str, optional, default: None
1039
+ Role to use for fetching secrets
1040
+ """
1041
+ ...
1042
+
1043
+ @typing.overload
1044
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1045
+ ...
1046
+
1047
+ @typing.overload
1048
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1049
+ ...
1050
+
1051
+ 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]]] = [], role: typing.Optional[str] = None):
1052
+ """
1053
+ Specifies secrets to be retrieved and injected as environment variables prior to
1054
+ the execution of a step.
1055
+
1056
+
1057
+ Parameters
1058
+ ----------
1059
+ sources : List[Union[str, Dict[str, Any]]], default: []
1060
+ List of secret specs, defining how the secrets are to be retrieved
1061
+ role : str, optional, default: None
1062
+ Role to use for fetching secrets
1063
+ """
1064
+ ...
1065
+
1066
+ @typing.overload
1067
+ 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]]]:
1068
+ """
1069
+ Specifies the resources needed when executing this step.
1070
+
1071
+ Use `@resources` to specify the resource requirements
1072
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
1073
+
1074
+ You can choose the compute layer on the command line by executing e.g.
1075
+ ```
1076
+ python myflow.py run --with batch
1077
+ ```
1078
+ or
1079
+ ```
1080
+ python myflow.py run --with kubernetes
1081
+ ```
1082
+ which executes the flow on the desired system using the
1083
+ requirements specified in `@resources`.
1084
+
1085
+
1086
+ Parameters
1087
+ ----------
1088
+ cpu : int, default 1
1089
+ Number of CPUs required for this step.
1090
+ gpu : int, optional, default None
1091
+ Number of GPUs required for this step.
1092
+ disk : int, optional, default None
1093
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
1094
+ memory : int, default 4096
1095
+ Memory size (in MB) required for this step.
1096
+ shared_memory : int, optional, default None
1097
+ The value for the size (in MiB) of the /dev/shm volume for this step.
1098
+ This parameter maps to the `--shm-size` option in Docker.
1099
+ """
1100
+ ...
1101
+
1102
+ @typing.overload
1103
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1104
+ ...
1105
+
1106
+ @typing.overload
1107
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1108
+ ...
1109
+
1110
+ 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):
1111
+ """
1112
+ Specifies the resources needed when executing this step.
1113
+
1114
+ Use `@resources` to specify the resource requirements
1115
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
1116
+
1117
+ You can choose the compute layer on the command line by executing e.g.
1118
+ ```
1119
+ python myflow.py run --with batch
1120
+ ```
1121
+ or
1122
+ ```
1123
+ python myflow.py run --with kubernetes
1124
+ ```
1125
+ which executes the flow on the desired system using the
1126
+ requirements specified in `@resources`.
1127
+
1128
+
1129
+ Parameters
1130
+ ----------
1131
+ cpu : int, default 1
1132
+ Number of CPUs required for this step.
1133
+ gpu : int, optional, default None
1134
+ Number of GPUs required for this step.
1135
+ disk : int, optional, default None
1136
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
1137
+ memory : int, default 4096
1138
+ Memory size (in MB) required for this step.
1139
+ shared_memory : int, optional, default None
1140
+ The value for the size (in MiB) of the /dev/shm volume for this step.
1141
+ This parameter maps to the `--shm-size` option in Docker.
1142
+ """
1143
+ ...
1144
+
1145
+ @typing.overload
1146
+ def app_deploy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1147
+ """
1148
+ Decorator prototype for all step decorators. This function gets specialized
1149
+ and imported for all decorators types by _import_plugin_decorators().
1150
+ """
1151
+ ...
1152
+
1153
+ @typing.overload
1154
+ def app_deploy(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1155
+ ...
1156
+
1157
+ def app_deploy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1158
+ """
1159
+ Decorator prototype for all step decorators. This function gets specialized
1160
+ and imported for all decorators types by _import_plugin_decorators().
1161
+ """
1162
+ ...
1163
+
1164
+ def s3_proxy(*, integration_name: typing.Optional[str] = None, write_mode: typing.Optional[str] = None, debug: typing.Optional[bool] = 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]]]:
1165
+ """
1166
+ S3 Proxy decorator for routing S3 requests through a local proxy service.
1167
+
1168
+
1169
+ Parameters
1170
+ ----------
1171
+ integration_name : str, optional
1172
+ Name of the S3 proxy integration. If not specified, will use the only
1173
+ available S3 proxy integration in the namespace (fails if multiple exist).
1174
+ write_mode : str, optional
1175
+ The desired behavior during write operations to target (origin) S3 bucket.
1176
+ allowed options are:
1177
+ "origin-and-cache" -> write to both the target S3 bucket and local object
1178
+ storage
1179
+ "origin" -> only write to the target S3 bucket
1180
+ "cache" -> only write to the object storage service used for caching
1181
+ debug : bool, optional
1182
+ Enable debug logging for proxy operations.
1183
+ """
1184
+ ...
1185
+
1186
+ @typing.overload
1187
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1188
+ """
1189
+ Internal decorator to support Fast bakery
1190
+ """
1191
+ ...
1192
+
1193
+ @typing.overload
1194
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1195
+ ...
1196
+
1197
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1198
+ """
1199
+ Internal decorator to support Fast bakery
1200
+ """
1201
+ ...
1202
+
1203
+ @typing.overload
1204
+ 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]]]:
1205
+ """
1206
+ Creates a human-readable report, a Metaflow Card, after this step completes.
1207
+
1208
+ Note that you may add multiple `@card` decorators in a step with different parameters.
1209
+
1210
+
1211
+ Parameters
1212
+ ----------
1213
+ type : str, default 'default'
1214
+ Card type.
1215
+ id : str, optional, default None
1216
+ If multiple cards are present, use this id to identify this card.
1217
+ options : Dict[str, Any], default {}
1218
+ Options passed to the card. The contents depend on the card type.
1219
+ timeout : int, default 45
1220
+ Interrupt reporting if it takes more than this many seconds.
1221
+ """
1222
+ ...
1223
+
1224
+ @typing.overload
1225
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1226
+ ...
1227
+
1228
+ @typing.overload
1229
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1230
+ ...
1231
+
1232
+ 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):
1233
+ """
1234
+ Creates a human-readable report, a Metaflow Card, after this step completes.
1235
+
1236
+ Note that you may add multiple `@card` decorators in a step with different parameters.
1237
+
1238
+
1239
+ Parameters
1240
+ ----------
1241
+ type : str, default 'default'
1242
+ Card type.
1243
+ id : str, optional, default None
1244
+ If multiple cards are present, use this id to identify this card.
1245
+ options : Dict[str, Any], default {}
1246
+ Options passed to the card. The contents depend on the card type.
1247
+ timeout : int, default 45
1248
+ Interrupt reporting if it takes more than this many seconds.
1249
+ """
1250
+ ...
1251
+
1252
+ @typing.overload
1253
+ def coreweave_s3_proxy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1254
+ """
1255
+ CoreWeave-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
1256
+ It exists to make it easier for users to know that this decorator should only be used with
1257
+ a Neo Cloud like CoreWeave.
1258
+ """
1259
+ ...
1260
+
1261
+ @typing.overload
1262
+ def coreweave_s3_proxy(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1263
+ ...
1264
+
1265
+ def coreweave_s3_proxy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1266
+ """
1267
+ CoreWeave-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
1268
+ It exists to make it easier for users to know that this decorator should only be used with
1269
+ a Neo Cloud like CoreWeave.
1270
+ """
1271
+ ...
1272
+
1273
+ @typing.overload
1274
+ 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]]]:
1275
+ """
1276
+ Specifies the Conda environment for the step.
1277
+
1278
+ Information in this decorator will augment any
1279
+ attributes set in the `@conda_base` flow-level decorator. Hence,
1280
+ you can use `@conda_base` to set packages required by all
1281
+ steps and use `@conda` to specify step-specific overrides.
1282
+
1283
+
1284
+ Parameters
1285
+ ----------
1286
+ packages : Dict[str, str], default {}
1287
+ Packages to use for this step. The key is the name of the package
1288
+ and the value is the version to use.
1289
+ libraries : Dict[str, str], default {}
1290
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1291
+ python : str, optional, default None
1292
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1293
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1294
+ disabled : bool, default False
1295
+ If set to True, disables @conda.
1296
+ """
1297
+ ...
1298
+
1299
+ @typing.overload
1300
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1301
+ ...
1302
+
1303
+ @typing.overload
1304
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1305
+ ...
1306
+
1307
+ 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):
1308
+ """
1309
+ Specifies the Conda environment for the step.
1310
+
1311
+ Information in this decorator will augment any
1312
+ attributes set in the `@conda_base` flow-level decorator. Hence,
1313
+ you can use `@conda_base` to set packages required by all
1314
+ steps and use `@conda` to specify step-specific overrides.
1315
+
1316
+
1317
+ Parameters
1318
+ ----------
1319
+ packages : Dict[str, str], default {}
1320
+ Packages to use for this step. The key is the name of the package
1321
+ and the value is the version to use.
1322
+ libraries : Dict[str, str], default {}
1323
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1324
+ python : str, optional, default None
1325
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1326
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1327
+ disabled : bool, default False
1328
+ If set to True, disables @conda.
1329
+ """
1330
+ ...
1331
+
1332
+ @typing.overload
1333
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1334
+ """
1335
+ Decorator prototype for all step decorators. This function gets specialized
1336
+ and imported for all decorators types by _import_plugin_decorators().
1337
+ """
1338
+ ...
1339
+
1340
+ @typing.overload
1341
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1342
+ ...
1343
+
1344
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1345
+ """
1346
+ Decorator prototype for all step decorators. This function gets specialized
1347
+ and imported for all decorators types by _import_plugin_decorators().
1348
+ """
1349
+ ...
1350
+
1351
+ def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None):
1352
+ """
1353
+ Allows setting external datastores to save data for the
1354
+ `@checkpoint`/`@model`/`@huggingface_hub` decorators.
1355
+
1356
+ This decorator is useful when users wish to save data to a different datastore
1357
+ than what is configured in Metaflow. This can be for variety of reasons:
1358
+
1359
+ 1. Data security: The objects needs to be stored in a bucket (object storage) that is not accessible by other flows.
1360
+ 2. Data Locality: The location where the task is executing is not located in the same region as the datastore.
1361
+ - Example: Metaflow datastore lives in US East, but the task is executing in Finland datacenters.
1362
+ 3. Data Lifecycle Policies: The objects need to be archived / managed separately from the Metaflow managed objects.
1363
+ - Example: Flow is training very large models that need to be stored separately and will be deleted more aggressively than the Metaflow managed objects.
1364
+
1365
+ Usage:
1366
+ ----------
1367
+
1368
+ - Using a custom IAM role to access the datastore.
1369
+
1370
+ ```python
1371
+ @with_artifact_store(
1372
+ type="s3",
1373
+ config=lambda: {
1374
+ "root": "s3://my-bucket-foo/path/to/root",
1375
+ "role_arn": ROLE,
1376
+ },
1377
+ )
1378
+ class MyFlow(FlowSpec):
1379
+
1380
+ @checkpoint
1381
+ @step
1382
+ def start(self):
1383
+ with open("my_file.txt", "w") as f:
1384
+ f.write("Hello, World!")
1385
+ self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1386
+ self.next(self.end)
1387
+
1388
+ ```
1389
+
1390
+ - Using credentials to access the s3-compatible datastore.
1391
+
1392
+ ```python
1393
+ @with_artifact_store(
1394
+ type="s3",
1395
+ config=lambda: {
1396
+ "root": "s3://my-bucket-foo/path/to/root",
1397
+ "client_params": {
1398
+ "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1399
+ "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1400
+ },
1401
+ },
1402
+ )
1403
+ class MyFlow(FlowSpec):
1404
+
1405
+ @checkpoint
1406
+ @step
1407
+ def start(self):
1408
+ with open("my_file.txt", "w") as f:
1409
+ f.write("Hello, World!")
1410
+ self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1411
+ self.next(self.end)
1471
1412
 
1472
- Use `@pypi_base` to set common packages required by all
1473
- steps and use `@pypi` to specify step-specific overrides.
1413
+ ```
1474
1414
 
1475
- Parameters
1476
- ----------
1477
- packages : Dict[str, str], default: {}
1478
- Packages to use for this flow. The key is the name of the package
1479
- and the value is the version to use.
1480
- python : str, optional, default: None
1481
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1482
- that the version used will correspond to the version of the Python interpreter used to start the run.
1483
- """
1484
- ...
1485
-
1486
- 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]]:
1487
- """
1488
- 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.
1489
- 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.
1415
+ - Accessing objects stored in external datastores after task execution.
1490
1416
 
1417
+ ```python
1418
+ run = Run("CheckpointsTestsFlow/8992")
1419
+ with artifact_store_from(run=run, config={
1420
+ "client_params": {
1421
+ "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1422
+ "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1423
+ },
1424
+ }):
1425
+ with Checkpoint() as cp:
1426
+ latest = cp.list(
1427
+ task=run["start"].task
1428
+ )[0]
1429
+ print(latest)
1430
+ cp.load(
1431
+ latest,
1432
+ "test-checkpoints"
1433
+ )
1491
1434
 
1492
- Parameters
1435
+ task = Task("TorchTuneFlow/8484/train/53673")
1436
+ with artifact_store_from(run=run, config={
1437
+ "client_params": {
1438
+ "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1439
+ "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1440
+ },
1441
+ }):
1442
+ load_model(
1443
+ task.data.model_ref,
1444
+ "test-models"
1445
+ )
1446
+ ```
1447
+ Parameters:
1493
1448
  ----------
1494
- timeout : int
1495
- Time, in seconds before the task times out and fails. (Default: 3600)
1496
- poke_interval : int
1497
- Time in seconds that the job should wait in between each try. (Default: 60)
1498
- mode : str
1499
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1500
- exponential_backoff : bool
1501
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1502
- pool : str
1503
- the slot pool this task should run in,
1504
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1505
- soft_fail : bool
1506
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1507
- name : str
1508
- Name of the sensor on Airflow
1509
- description : str
1510
- Description of sensor in the Airflow UI
1511
- external_dag_id : str
1512
- The dag_id that contains the task you want to wait for.
1513
- external_task_ids : List[str]
1514
- The list of task_ids that you want to wait for.
1515
- If None (default value) the sensor waits for the DAG. (Default: None)
1516
- allowed_states : List[str]
1517
- Iterable of allowed states, (Default: ['success'])
1518
- failed_states : List[str]
1519
- Iterable of failed or dis-allowed states. (Default: None)
1520
- execution_delta : datetime.timedelta
1521
- time difference with the previous execution to look at,
1522
- the default is the same logical date as the current task or DAG. (Default: None)
1523
- check_existence: bool
1524
- Set to True to check if the external task exists or check if
1525
- the DAG to wait for exists. (Default: True)
1449
+
1450
+ type: str
1451
+ The type of the datastore. Can be one of 's3', 'gcs', 'azure' or any other supported metaflow Datastore.
1452
+
1453
+ config: dict or Callable
1454
+ Dictionary of configuration options for the datastore. The following keys are required:
1455
+ - root: The root path in the datastore where the data will be saved. (needs to be in the format expected by the datastore)
1456
+ - example: 's3://bucket-name/path/to/root'
1457
+ - example: 'gs://bucket-name/path/to/root'
1458
+ - example: 'https://myblockacc.blob.core.windows.net/metaflow/'
1459
+ - role_arn (optional): AWS IAM role to access s3 bucket (only when `type` is 's3')
1460
+ - session_vars (optional): AWS session variables to access s3 bucket (only when `type` is 's3')
1461
+ - client_params (optional): AWS client parameters to access s3 bucket (only when `type` is 's3')
1526
1462
  """
1527
1463
  ...
1528
1464
 
1529
- def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1465
+ @typing.overload
1466
+ 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]]:
1530
1467
  """
1531
- Specifies what flows belong to the same project.
1468
+ Specifies the flow(s) that this flow depends on.
1532
1469
 
1533
- A project-specific namespace is created for all flows that
1534
- use the same `@project(name)`.
1470
+ ```
1471
+ @trigger_on_finish(flow='FooFlow')
1472
+ ```
1473
+ or
1474
+ ```
1475
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1476
+ ```
1477
+ This decorator respects the @project decorator and triggers the flow
1478
+ when upstream runs within the same namespace complete successfully
1535
1479
 
1480
+ Additionally, you can specify project aware upstream flow dependencies
1481
+ by specifying the fully qualified project_flow_name.
1482
+ ```
1483
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1484
+ ```
1485
+ or
1486
+ ```
1487
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1488
+ ```
1536
1489
 
1537
- Parameters
1538
- ----------
1539
- name : str
1540
- Project name. Make sure that the name is unique amongst all
1541
- projects that use the same production scheduler. The name may
1542
- contain only lowercase alphanumeric characters and underscores.
1490
+ You can also specify just the project or project branch (other values will be
1491
+ inferred from the current project or project branch):
1492
+ ```
1493
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1494
+ ```
1543
1495
 
1544
- branch : Optional[str], default None
1545
- The branch to use. If not specified, the branch is set to
1546
- `user.<username>` unless `production` is set to `True`. This can
1547
- also be set on the command line using `--branch` as a top-level option.
1548
- It is an error to specify `branch` in the decorator and on the command line.
1496
+ Note that `branch` is typically one of:
1497
+ - `prod`
1498
+ - `user.bob`
1499
+ - `test.my_experiment`
1500
+ - `prod.staging`
1549
1501
 
1550
- production : bool, default False
1551
- Whether or not the branch is the production branch. This can also be set on the
1552
- command line using `--production` as a top-level option. It is an error to specify
1553
- `production` in the decorator and on the command line.
1554
- The project branch name will be:
1555
- - if `branch` is specified:
1556
- - if `production` is True: `prod.<branch>`
1557
- - if `production` is False: `test.<branch>`
1558
- - if `branch` is not specified:
1559
- - if `production` is True: `prod`
1560
- - if `production` is False: `user.<username>`
1502
+
1503
+ Parameters
1504
+ ----------
1505
+ flow : Union[str, Dict[str, str]], optional, default None
1506
+ Upstream flow dependency for this flow.
1507
+ flows : List[Union[str, Dict[str, str]]], default []
1508
+ Upstream flow dependencies for this flow.
1509
+ options : Dict[str, Any], default {}
1510
+ Backend-specific configuration for tuning eventing behavior.
1561
1511
  """
1562
1512
  ...
1563
1513
 
1564
1514
  @typing.overload
1565
- 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]]:
1515
+ def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1516
+ ...
1517
+
1518
+ 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] = {}):
1566
1519
  """
1567
1520
  Specifies the flow(s) that this flow depends on.
1568
1521
 
@@ -1610,169 +1563,173 @@ def trigger_on_finish(*, flow: typing.Union[typing.Dict[str, str], str, None] =
1610
1563
  """
1611
1564
  ...
1612
1565
 
1566
+ def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1567
+ """
1568
+ Specifies what flows belong to the same project.
1569
+
1570
+ A project-specific namespace is created for all flows that
1571
+ use the same `@project(name)`.
1572
+
1573
+
1574
+ Parameters
1575
+ ----------
1576
+ name : str
1577
+ Project name. Make sure that the name is unique amongst all
1578
+ projects that use the same production scheduler. The name may
1579
+ contain only lowercase alphanumeric characters and underscores.
1580
+
1581
+ branch : Optional[str], default None
1582
+ The branch to use. If not specified, the branch is set to
1583
+ `user.<username>` unless `production` is set to `True`. This can
1584
+ also be set on the command line using `--branch` as a top-level option.
1585
+ It is an error to specify `branch` in the decorator and on the command line.
1586
+
1587
+ production : bool, default False
1588
+ Whether or not the branch is the production branch. This can also be set on the
1589
+ command line using `--production` as a top-level option. It is an error to specify
1590
+ `production` in the decorator and on the command line.
1591
+ The project branch name will be:
1592
+ - if `branch` is specified:
1593
+ - if `production` is True: `prod.<branch>`
1594
+ - if `production` is False: `test.<branch>`
1595
+ - if `branch` is not specified:
1596
+ - if `production` is True: `prod`
1597
+ - if `production` is False: `user.<username>`
1598
+ """
1599
+ ...
1600
+
1601
+ @typing.overload
1602
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1603
+ """
1604
+ Specifies the PyPI packages for all steps of the flow.
1605
+
1606
+ Use `@pypi_base` to set common packages required by all
1607
+ steps and use `@pypi` to specify step-specific overrides.
1608
+
1609
+ Parameters
1610
+ ----------
1611
+ packages : Dict[str, str], default: {}
1612
+ Packages to use for this flow. The key is the name of the package
1613
+ and the value is the version to use.
1614
+ python : str, optional, default: None
1615
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1616
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1617
+ """
1618
+ ...
1619
+
1620
+ @typing.overload
1621
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1622
+ ...
1623
+
1624
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1625
+ """
1626
+ Specifies the PyPI packages for all steps of the flow.
1627
+
1628
+ Use `@pypi_base` to set common packages required by all
1629
+ steps and use `@pypi` to specify step-specific overrides.
1630
+
1631
+ Parameters
1632
+ ----------
1633
+ packages : Dict[str, str], default: {}
1634
+ Packages to use for this flow. The key is the name of the package
1635
+ and the value is the version to use.
1636
+ python : str, optional, default: None
1637
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1638
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1639
+ """
1640
+ ...
1641
+
1642
+ @typing.overload
1643
+ def schedule(*, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1644
+ """
1645
+ Specifies the times when the flow should be run when running on a
1646
+ production scheduler.
1647
+
1648
+
1649
+ Parameters
1650
+ ----------
1651
+ hourly : bool, default False
1652
+ Run the workflow hourly.
1653
+ daily : bool, default True
1654
+ Run the workflow daily.
1655
+ weekly : bool, default False
1656
+ Run the workflow weekly.
1657
+ cron : str, optional, default None
1658
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1659
+ specified by this expression.
1660
+ timezone : str, optional, default None
1661
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1662
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1663
+ """
1664
+ ...
1665
+
1613
1666
  @typing.overload
1614
- def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1667
+ def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1615
1668
  ...
1616
1669
 
1617
- 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] = {}):
1670
+ def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None):
1618
1671
  """
1619
- Specifies the flow(s) that this flow depends on.
1620
-
1621
- ```
1622
- @trigger_on_finish(flow='FooFlow')
1623
- ```
1624
- or
1625
- ```
1626
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1627
- ```
1628
- This decorator respects the @project decorator and triggers the flow
1629
- when upstream runs within the same namespace complete successfully
1630
-
1631
- Additionally, you can specify project aware upstream flow dependencies
1632
- by specifying the fully qualified project_flow_name.
1633
- ```
1634
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1635
- ```
1636
- or
1637
- ```
1638
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1639
- ```
1640
-
1641
- You can also specify just the project or project branch (other values will be
1642
- inferred from the current project or project branch):
1643
- ```
1644
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1645
- ```
1646
-
1647
- Note that `branch` is typically one of:
1648
- - `prod`
1649
- - `user.bob`
1650
- - `test.my_experiment`
1651
- - `prod.staging`
1672
+ Specifies the times when the flow should be run when running on a
1673
+ production scheduler.
1652
1674
 
1653
1675
 
1654
1676
  Parameters
1655
1677
  ----------
1656
- flow : Union[str, Dict[str, str]], optional, default None
1657
- Upstream flow dependency for this flow.
1658
- flows : List[Union[str, Dict[str, str]]], default []
1659
- Upstream flow dependencies for this flow.
1660
- options : Dict[str, Any], default {}
1661
- Backend-specific configuration for tuning eventing behavior.
1678
+ hourly : bool, default False
1679
+ Run the workflow hourly.
1680
+ daily : bool, default True
1681
+ Run the workflow daily.
1682
+ weekly : bool, default False
1683
+ Run the workflow weekly.
1684
+ cron : str, optional, default None
1685
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1686
+ specified by this expression.
1687
+ timezone : str, optional, default None
1688
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1689
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1662
1690
  """
1663
1691
  ...
1664
1692
 
1665
- def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None):
1693
+ 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]]:
1666
1694
  """
1667
- Allows setting external datastores to save data for the
1668
- `@checkpoint`/`@model`/`@huggingface_hub` decorators.
1669
-
1670
- This decorator is useful when users wish to save data to a different datastore
1671
- than what is configured in Metaflow. This can be for variety of reasons:
1672
-
1673
- 1. Data security: The objects needs to be stored in a bucket (object storage) that is not accessible by other flows.
1674
- 2. Data Locality: The location where the task is executing is not located in the same region as the datastore.
1675
- - Example: Metaflow datastore lives in US East, but the task is executing in Finland datacenters.
1676
- 3. Data Lifecycle Policies: The objects need to be archived / managed separately from the Metaflow managed objects.
1677
- - Example: Flow is training very large models that need to be stored separately and will be deleted more aggressively than the Metaflow managed objects.
1678
-
1679
- Usage:
1680
- ----------
1681
-
1682
- - Using a custom IAM role to access the datastore.
1683
-
1684
- ```python
1685
- @with_artifact_store(
1686
- type="s3",
1687
- config=lambda: {
1688
- "root": "s3://my-bucket-foo/path/to/root",
1689
- "role_arn": ROLE,
1690
- },
1691
- )
1692
- class MyFlow(FlowSpec):
1693
-
1694
- @checkpoint
1695
- @step
1696
- def start(self):
1697
- with open("my_file.txt", "w") as f:
1698
- f.write("Hello, World!")
1699
- self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1700
- self.next(self.end)
1701
-
1702
- ```
1703
-
1704
- - Using credentials to access the s3-compatible datastore.
1705
-
1706
- ```python
1707
- @with_artifact_store(
1708
- type="s3",
1709
- config=lambda: {
1710
- "root": "s3://my-bucket-foo/path/to/root",
1711
- "client_params": {
1712
- "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1713
- "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1714
- },
1715
- },
1716
- )
1717
- class MyFlow(FlowSpec):
1718
-
1719
- @checkpoint
1720
- @step
1721
- def start(self):
1722
- with open("my_file.txt", "w") as f:
1723
- f.write("Hello, World!")
1724
- self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1725
- self.next(self.end)
1726
-
1727
- ```
1728
-
1729
- - Accessing objects stored in external datastores after task execution.
1695
+ 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.
1696
+ 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.
1730
1697
 
1731
- ```python
1732
- run = Run("CheckpointsTestsFlow/8992")
1733
- with artifact_store_from(run=run, config={
1734
- "client_params": {
1735
- "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1736
- "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1737
- },
1738
- }):
1739
- with Checkpoint() as cp:
1740
- latest = cp.list(
1741
- task=run["start"].task
1742
- )[0]
1743
- print(latest)
1744
- cp.load(
1745
- latest,
1746
- "test-checkpoints"
1747
- )
1748
1698
 
1749
- task = Task("TorchTuneFlow/8484/train/53673")
1750
- with artifact_store_from(run=run, config={
1751
- "client_params": {
1752
- "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1753
- "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1754
- },
1755
- }):
1756
- load_model(
1757
- task.data.model_ref,
1758
- "test-models"
1759
- )
1760
- ```
1761
- Parameters:
1699
+ Parameters
1762
1700
  ----------
1763
-
1764
- type: str
1765
- The type of the datastore. Can be one of 's3', 'gcs', 'azure' or any other supported metaflow Datastore.
1766
-
1767
- config: dict or Callable
1768
- Dictionary of configuration options for the datastore. The following keys are required:
1769
- - root: The root path in the datastore where the data will be saved. (needs to be in the format expected by the datastore)
1770
- - example: 's3://bucket-name/path/to/root'
1771
- - example: 'gs://bucket-name/path/to/root'
1772
- - example: 'https://myblockacc.blob.core.windows.net/metaflow/'
1773
- - role_arn (optional): AWS IAM role to access s3 bucket (only when `type` is 's3')
1774
- - session_vars (optional): AWS session variables to access s3 bucket (only when `type` is 's3')
1775
- - client_params (optional): AWS client parameters to access s3 bucket (only when `type` is 's3')
1701
+ timeout : int
1702
+ Time, in seconds before the task times out and fails. (Default: 3600)
1703
+ poke_interval : int
1704
+ Time in seconds that the job should wait in between each try. (Default: 60)
1705
+ mode : str
1706
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1707
+ exponential_backoff : bool
1708
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1709
+ pool : str
1710
+ the slot pool this task should run in,
1711
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1712
+ soft_fail : bool
1713
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1714
+ name : str
1715
+ Name of the sensor on Airflow
1716
+ description : str
1717
+ Description of sensor in the Airflow UI
1718
+ external_dag_id : str
1719
+ The dag_id that contains the task you want to wait for.
1720
+ external_task_ids : List[str]
1721
+ The list of task_ids that you want to wait for.
1722
+ If None (default value) the sensor waits for the DAG. (Default: None)
1723
+ allowed_states : List[str]
1724
+ Iterable of allowed states, (Default: ['success'])
1725
+ failed_states : List[str]
1726
+ Iterable of failed or dis-allowed states. (Default: None)
1727
+ execution_delta : datetime.timedelta
1728
+ time difference with the previous execution to look at,
1729
+ the default is the same logical date as the current task or DAG. (Default: None)
1730
+ check_existence: bool
1731
+ Set to True to check if the external task exists or check if
1732
+ the DAG to wait for exists. (Default: True)
1776
1733
  """
1777
1734
  ...
1778
1735
 
@@ -1827,6 +1784,49 @@ def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packa
1827
1784
  """
1828
1785
  ...
1829
1786
 
1787
+ 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]]:
1788
+ """
1789
+ 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)
1790
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1791
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1792
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
1793
+ starts only after all sensors finish.
1794
+
1795
+
1796
+ Parameters
1797
+ ----------
1798
+ timeout : int
1799
+ Time, in seconds before the task times out and fails. (Default: 3600)
1800
+ poke_interval : int
1801
+ Time in seconds that the job should wait in between each try. (Default: 60)
1802
+ mode : str
1803
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1804
+ exponential_backoff : bool
1805
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1806
+ pool : str
1807
+ the slot pool this task should run in,
1808
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1809
+ soft_fail : bool
1810
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1811
+ name : str
1812
+ Name of the sensor on Airflow
1813
+ description : str
1814
+ Description of sensor in the Airflow UI
1815
+ bucket_key : Union[str, List[str]]
1816
+ The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1817
+ When it's specified as a full s3:// url, please leave `bucket_name` as None
1818
+ bucket_name : str
1819
+ Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1820
+ When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1821
+ wildcard_match : bool
1822
+ whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1823
+ aws_conn_id : str
1824
+ a reference to the s3 connection on Airflow. (Default: None)
1825
+ verify : bool
1826
+ Whether or not to verify SSL certificates for S3 connection. (Default: None)
1827
+ """
1828
+ ...
1829
+
1830
1830
  @typing.overload
1831
1831
  def trigger(*, event: typing.Union[str, typing.Dict[str, typing.Any], None] = None, events: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], options: typing.Dict[str, typing.Any] = {}) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1832
1832
  """