ob-metaflow-stubs 6.0.10.4__py2.py3-none-any.whl → 6.0.10.6__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 +766 -766
  2. metaflow-stubs/cards.pyi +1 -1
  3. metaflow-stubs/cli.pyi +1 -1
  4. metaflow-stubs/cli_components/__init__.pyi +1 -1
  5. metaflow-stubs/cli_components/utils.pyi +1 -1
  6. metaflow-stubs/client/__init__.pyi +1 -1
  7. metaflow-stubs/client/core.pyi +4 -4
  8. metaflow-stubs/client/filecache.pyi +2 -2
  9. metaflow-stubs/events.pyi +1 -1
  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 +73 -73
  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 +1 -1
  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 +2 -2
  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 +3 -3
  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 +4 -4
  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 +2 -2
  37. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/decorator.pyi +2 -2
  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 +1 -1
  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 +3 -3
  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 +1 -1
  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 +2 -2
  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 +1 -1
  73. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/app_config.pyi +2 -2
  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 +3 -3
  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 +4 -4
  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 +6 -6
  119. metaflow-stubs/packaging_sys/backend.pyi +4 -4
  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 +2 -2
  125. metaflow-stubs/plugins/__init__.pyi +9 -9
  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 +1 -1
  138. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +3 -3
  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 +2 -2
  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 +3 -3
  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 +2 -2
  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 +1 -1
  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 +2 -2
  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 +1 -1
  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 +2 -2
  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 +2 -2
  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 +2 -2
  213. metaflow-stubs/plugins/pypi/conda_decorator.pyi +1 -1
  214. metaflow-stubs/plugins/pypi/conda_environment.pyi +4 -4
  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 +2 -2
  222. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +1 -1
  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 +1 -1
  230. metaflow-stubs/plugins/timeout_decorator.pyi +1 -1
  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 +33 -33
  238. metaflow-stubs/runner/deployer_impl.pyi +1 -1
  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 +1 -1
  244. metaflow-stubs/system/__init__.pyi +1 -1
  245. metaflow-stubs/system/system_logger.pyi +1 -1
  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 +5 -5
  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 +4 -4
  255. metaflow-stubs/user_decorators/mutable_step.pyi +3 -3
  256. metaflow-stubs/user_decorators/user_flow_decorator.pyi +4 -4
  257. metaflow-stubs/user_decorators/user_step_decorator.pyi +5 -5
  258. {ob_metaflow_stubs-6.0.10.4.dist-info → ob_metaflow_stubs-6.0.10.6.dist-info}/METADATA +1 -1
  259. ob_metaflow_stubs-6.0.10.6.dist-info/RECORD +262 -0
  260. ob_metaflow_stubs-6.0.10.4.dist-info/RECORD +0 -262
  261. {ob_metaflow_stubs-6.0.10.4.dist-info → ob_metaflow_stubs-6.0.10.6.dist-info}/WHEEL +0 -0
  262. {ob_metaflow_stubs-6.0.10.4.dist-info → ob_metaflow_stubs-6.0.10.6.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.18.5.1+obcheckpoint(0.2.6);ob(v1) #
4
- # Generated on 2025-09-16T23:23:08.891416 #
4
+ # Generated on 2025-09-19T18:02:17.702045 #
5
5
  ######################################################################################################
6
6
 
7
7
  from __future__ import annotations
@@ -48,9 +48,9 @@ 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 conda_environment_yml_parser as conda_environment_yml_parser
52
51
  from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
53
52
  from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
53
+ from .plugins.pypi.parsers import conda_environment_yml_parser as conda_environment_yml_parser
54
54
  from . import client as client
55
55
  from .client.core import namespace as namespace
56
56
  from .client.core import get_namespace as get_namespace
@@ -167,6 +167,104 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
167
167
  """
168
168
  ...
169
169
 
170
+ 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]]]:
171
+ """
172
+ Specifies that this step should execute on DGX cloud.
173
+
174
+
175
+ Parameters
176
+ ----------
177
+ gpu : int
178
+ Number of GPUs to use.
179
+ gpu_type : str
180
+ Type of Nvidia GPU to use.
181
+ """
182
+ ...
183
+
184
+ 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]]]:
185
+ """
186
+ This decorator is used to run Ollama APIs as Metaflow task sidecars.
187
+
188
+ User code call
189
+ --------------
190
+ @ollama(
191
+ models=[...],
192
+ ...
193
+ )
194
+
195
+ Valid backend options
196
+ ---------------------
197
+ - 'local': Run as a separate process on the local task machine.
198
+ - (TODO) 'managed': Outerbounds hosts and selects compute provider.
199
+ - (TODO) 'remote': Spin up separate instance to serve Ollama models.
200
+
201
+ Valid model options
202
+ -------------------
203
+ Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
204
+
205
+
206
+ Parameters
207
+ ----------
208
+ models: list[str]
209
+ List of Ollama containers running models in sidecars.
210
+ backend: str
211
+ Determines where and how to run the Ollama process.
212
+ force_pull: bool
213
+ Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
214
+ cache_update_policy: str
215
+ Cache update policy: "auto", "force", or "never".
216
+ force_cache_update: bool
217
+ Simple override for "force" cache update policy.
218
+ debug: bool
219
+ Whether to turn on verbose debugging logs.
220
+ circuit_breaker_config: dict
221
+ Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
222
+ timeout_config: dict
223
+ Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
224
+ """
225
+ ...
226
+
227
+ @typing.overload
228
+ def test_append_card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
229
+ """
230
+ A simple decorator that demonstrates using CardDecoratorInjector
231
+ to inject a card and render simple markdown content.
232
+ """
233
+ ...
234
+
235
+ @typing.overload
236
+ def test_append_card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
237
+ ...
238
+
239
+ def test_append_card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
240
+ """
241
+ A simple decorator that demonstrates using CardDecoratorInjector
242
+ to inject a card and render simple markdown content.
243
+ """
244
+ ...
245
+
246
+ 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]]]:
247
+ """
248
+ S3 Proxy decorator for routing S3 requests through a local proxy service.
249
+
250
+
251
+ Parameters
252
+ ----------
253
+ integration_name : str, optional
254
+ Name of the S3 proxy integration. If not specified, will use the only
255
+ available S3 proxy integration in the namespace (fails if multiple exist).
256
+ write_mode : str, optional
257
+ The desired behavior during write operations to target (origin) S3 bucket.
258
+ allowed options are:
259
+ "origin-and-cache" -> write to both the target S3 bucket and local object
260
+ storage
261
+ "origin" -> only write to the target S3 bucket
262
+ "cache" -> only write to the object storage service used for caching
263
+ debug : bool, optional
264
+ Enable debug logging for proxy operations.
265
+ """
266
+ ...
267
+
170
268
  @typing.overload
171
269
  def coreweave_s3_proxy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
172
270
  """
@@ -205,25 +303,6 @@ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepF
205
303
  """
206
304
  ...
207
305
 
208
- @typing.overload
209
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
210
- """
211
- Decorator prototype for all step decorators. This function gets specialized
212
- and imported for all decorators types by _import_plugin_decorators().
213
- """
214
- ...
215
-
216
- @typing.overload
217
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
218
- ...
219
-
220
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
221
- """
222
- Decorator prototype for all step decorators. This function gets specialized
223
- and imported for all decorators types by _import_plugin_decorators().
224
- """
225
- ...
226
-
227
306
  @typing.overload
228
307
  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]]]:
229
308
  """
@@ -279,55 +358,6 @@ def retry(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
279
358
  """
280
359
  ...
281
360
 
282
- @typing.overload
283
- 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]]]:
284
- """
285
- Creates a human-readable report, a Metaflow Card, after this step completes.
286
-
287
- Note that you may add multiple `@card` decorators in a step with different parameters.
288
-
289
-
290
- Parameters
291
- ----------
292
- type : str, default 'default'
293
- Card type.
294
- id : str, optional, default None
295
- If multiple cards are present, use this id to identify this card.
296
- options : Dict[str, Any], default {}
297
- Options passed to the card. The contents depend on the card type.
298
- timeout : int, default 45
299
- Interrupt reporting if it takes more than this many seconds.
300
- """
301
- ...
302
-
303
- @typing.overload
304
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
305
- ...
306
-
307
- @typing.overload
308
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
309
- ...
310
-
311
- 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):
312
- """
313
- Creates a human-readable report, a Metaflow Card, after this step completes.
314
-
315
- Note that you may add multiple `@card` decorators in a step with different parameters.
316
-
317
-
318
- Parameters
319
- ----------
320
- type : str, default 'default'
321
- Card type.
322
- id : str, optional, default None
323
- If multiple cards are present, use this id to identify this card.
324
- options : Dict[str, Any], default {}
325
- Options passed to the card. The contents depend on the card type.
326
- timeout : int, default 45
327
- Interrupt reporting if it takes more than this many seconds.
328
- """
329
- ...
330
-
331
361
  def huggingface_hub(*, temp_dir_root: typing.Optional[str] = None, cache_scope: 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]]]:
332
362
  """
333
363
  Decorator that helps cache, version, and store models/datasets from the Hugging Face Hub.
@@ -455,190 +485,196 @@ def huggingface_hub(*, temp_dir_root: typing.Optional[str] = None, cache_scope:
455
485
  """
456
486
  ...
457
487
 
458
- 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]]]:
488
+ @typing.overload
489
+ 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]]]:
459
490
  """
460
- This decorator is used to run vllm APIs as Metaflow task sidecars.
461
-
462
- User code call
463
- --------------
464
- @vllm(
465
- model="...",
466
- ...
467
- )
468
-
469
- Valid backend options
470
- ---------------------
471
- - 'local': Run as a separate process on the local task machine.
472
-
473
- Valid model options
474
- -------------------
475
- Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
476
-
477
- NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
478
- If you need multiple models, you must create multiple @vllm decorators.
491
+ Specifies the resources needed when executing this step.
479
492
 
493
+ Use `@resources` to specify the resource requirements
494
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
480
495
 
481
- Parameters
482
- ----------
483
- model: str
484
- HuggingFace model identifier to be served by vLLM.
485
- backend: str
486
- Determines where and how to run the vLLM process.
487
- openai_api_server: bool
488
- Whether to use OpenAI-compatible API server mode (subprocess) instead of native engine.
489
- Default is False (uses native engine).
490
- Set to True for backward compatibility with existing code.
491
- debug: bool
492
- Whether to turn on verbose debugging logs.
493
- card_refresh_interval: int
494
- Interval in seconds for refreshing the vLLM status card.
495
- Only used when openai_api_server=True.
496
- max_retries: int
497
- Maximum number of retries checking for vLLM server startup.
498
- Only used when openai_api_server=True.
499
- retry_alert_frequency: int
500
- Frequency of alert logs for vLLM server startup retries.
501
- Only used when openai_api_server=True.
502
- engine_args : dict
503
- Additional keyword arguments to pass to the vLLM engine.
504
- For example, `tensor_parallel_size=2`.
505
- """
506
- ...
507
-
508
- @typing.overload
509
- 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]]]:
510
- """
511
- Specifies secrets to be retrieved and injected as environment variables prior to
512
- the execution of a step.
496
+ You can choose the compute layer on the command line by executing e.g.
497
+ ```
498
+ python myflow.py run --with batch
499
+ ```
500
+ or
501
+ ```
502
+ python myflow.py run --with kubernetes
503
+ ```
504
+ which executes the flow on the desired system using the
505
+ requirements specified in `@resources`.
513
506
 
514
507
 
515
508
  Parameters
516
509
  ----------
517
- sources : List[Union[str, Dict[str, Any]]], default: []
518
- List of secret specs, defining how the secrets are to be retrieved
519
- role : str, optional, default: None
520
- Role to use for fetching secrets
510
+ cpu : int, default 1
511
+ Number of CPUs required for this step.
512
+ gpu : int, optional, default None
513
+ Number of GPUs required for this step.
514
+ disk : int, optional, default None
515
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
516
+ memory : int, default 4096
517
+ Memory size (in MB) required for this step.
518
+ shared_memory : int, optional, default None
519
+ The value for the size (in MiB) of the /dev/shm volume for this step.
520
+ This parameter maps to the `--shm-size` option in Docker.
521
521
  """
522
522
  ...
523
523
 
524
524
  @typing.overload
525
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
525
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
526
526
  ...
527
527
 
528
528
  @typing.overload
529
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
529
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
530
530
  ...
531
531
 
532
- 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):
532
+ 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):
533
533
  """
534
- Specifies secrets to be retrieved and injected as environment variables prior to
535
- the execution of a step.
534
+ Specifies the resources needed when executing this step.
535
+
536
+ Use `@resources` to specify the resource requirements
537
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
538
+
539
+ You can choose the compute layer on the command line by executing e.g.
540
+ ```
541
+ python myflow.py run --with batch
542
+ ```
543
+ or
544
+ ```
545
+ python myflow.py run --with kubernetes
546
+ ```
547
+ which executes the flow on the desired system using the
548
+ requirements specified in `@resources`.
536
549
 
537
550
 
538
551
  Parameters
539
552
  ----------
540
- sources : List[Union[str, Dict[str, Any]]], default: []
541
- List of secret specs, defining how the secrets are to be retrieved
542
- role : str, optional, default: None
543
- Role to use for fetching secrets
553
+ cpu : int, default 1
554
+ Number of CPUs required for this step.
555
+ gpu : int, optional, default None
556
+ Number of GPUs required for this step.
557
+ disk : int, optional, default None
558
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
559
+ memory : int, default 4096
560
+ Memory size (in MB) required for this step.
561
+ shared_memory : int, optional, default None
562
+ The value for the size (in MiB) of the /dev/shm volume for this step.
563
+ This parameter maps to the `--shm-size` option in Docker.
544
564
  """
545
565
  ...
546
566
 
547
- 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]]]:
567
+ 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]]]:
548
568
  """
549
- S3 Proxy decorator for routing S3 requests through a local proxy service.
569
+ Specifies that this step should execute on DGX cloud.
550
570
 
551
571
 
552
572
  Parameters
553
573
  ----------
554
- integration_name : str, optional
555
- Name of the S3 proxy integration. If not specified, will use the only
556
- available S3 proxy integration in the namespace (fails if multiple exist).
557
- write_mode : str, optional
558
- The desired behavior during write operations to target (origin) S3 bucket.
559
- allowed options are:
560
- "origin-and-cache" -> write to both the target S3 bucket and local object
561
- storage
562
- "origin" -> only write to the target S3 bucket
563
- "cache" -> only write to the object storage service used for caching
564
- debug : bool, optional
565
- Enable debug logging for proxy operations.
574
+ gpu : int
575
+ Number of GPUs to use.
576
+ gpu_type : str
577
+ Type of Nvidia GPU to use.
578
+ queue_timeout : int
579
+ Time to keep the job in NVCF's queue.
566
580
  """
567
581
  ...
568
582
 
569
- 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]]]:
583
+ @typing.overload
584
+ def timeout(*, seconds: int = 0, minutes: int = 0, hours: int = 0) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
570
585
  """
571
- This decorator is used to run Ollama APIs as Metaflow task sidecars.
586
+ Specifies a timeout for your step.
572
587
 
573
- User code call
574
- --------------
575
- @ollama(
576
- models=[...],
577
- ...
578
- )
588
+ This decorator is useful if this step may hang indefinitely.
579
589
 
580
- Valid backend options
581
- ---------------------
582
- - 'local': Run as a separate process on the local task machine.
583
- - (TODO) 'managed': Outerbounds hosts and selects compute provider.
584
- - (TODO) 'remote': Spin up separate instance to serve Ollama models.
590
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
591
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
592
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
585
593
 
586
- Valid model options
587
- -------------------
588
- Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
594
+ Note that all the values specified in parameters are added together so if you specify
595
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
589
596
 
590
597
 
591
598
  Parameters
592
599
  ----------
593
- models: list[str]
594
- List of Ollama containers running models in sidecars.
595
- backend: str
596
- Determines where and how to run the Ollama process.
597
- force_pull: bool
598
- Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
599
- cache_update_policy: str
600
- Cache update policy: "auto", "force", or "never".
601
- force_cache_update: bool
602
- Simple override for "force" cache update policy.
603
- debug: bool
604
- Whether to turn on verbose debugging logs.
605
- circuit_breaker_config: dict
606
- Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
607
- timeout_config: dict
608
- Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
600
+ seconds : int, default 0
601
+ Number of seconds to wait prior to timing out.
602
+ minutes : int, default 0
603
+ Number of minutes to wait prior to timing out.
604
+ hours : int, default 0
605
+ Number of hours to wait prior to timing out.
609
606
  """
610
607
  ...
611
608
 
612
- 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]]]:
609
+ @typing.overload
610
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
611
+ ...
612
+
613
+ @typing.overload
614
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
615
+ ...
616
+
617
+ def timeout(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, seconds: int = 0, minutes: int = 0, hours: int = 0):
613
618
  """
614
- Specifies that this step should execute on DGX cloud.
619
+ Specifies a timeout for your step.
620
+
621
+ This decorator is useful if this step may hang indefinitely.
622
+
623
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
624
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
625
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
626
+
627
+ Note that all the values specified in parameters are added together so if you specify
628
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
615
629
 
616
630
 
617
631
  Parameters
618
632
  ----------
619
- gpu : int
620
- Number of GPUs to use.
621
- gpu_type : str
622
- Type of Nvidia GPU to use.
633
+ seconds : int, default 0
634
+ Number of seconds to wait prior to timing out.
635
+ minutes : int, default 0
636
+ Number of minutes to wait prior to timing out.
637
+ hours : int, default 0
638
+ Number of hours to wait prior to timing out.
623
639
  """
624
640
  ...
625
641
 
626
642
  @typing.overload
627
- def app_deploy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
643
+ 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]]]:
628
644
  """
629
- Decorator prototype for all step decorators. This function gets specialized
630
- and imported for all decorators types by _import_plugin_decorators().
645
+ Specifies secrets to be retrieved and injected as environment variables prior to
646
+ the execution of a step.
647
+
648
+
649
+ Parameters
650
+ ----------
651
+ sources : List[Union[str, Dict[str, Any]]], default: []
652
+ List of secret specs, defining how the secrets are to be retrieved
653
+ role : str, optional, default: None
654
+ Role to use for fetching secrets
631
655
  """
632
656
  ...
633
657
 
634
658
  @typing.overload
635
- def app_deploy(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
659
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
636
660
  ...
637
661
 
638
- def app_deploy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
662
+ @typing.overload
663
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
664
+ ...
665
+
666
+ 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):
639
667
  """
640
- Decorator prototype for all step decorators. This function gets specialized
641
- and imported for all decorators types by _import_plugin_decorators().
668
+ Specifies secrets to be retrieved and injected as environment variables prior to
669
+ the execution of a step.
670
+
671
+
672
+ Parameters
673
+ ----------
674
+ sources : List[Union[str, Dict[str, Any]]], default: []
675
+ List of secret specs, defining how the secrets are to be retrieved
676
+ role : str, optional, default: None
677
+ Role to use for fetching secrets
642
678
  """
643
679
  ...
644
680
 
@@ -664,76 +700,273 @@ def nebius_s3_proxy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag],
664
700
  ...
665
701
 
666
702
  @typing.overload
667
- 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]]]:
703
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
668
704
  """
669
- Specifies the PyPI packages for the step.
705
+ Decorator prototype for all step decorators. This function gets specialized
706
+ and imported for all decorators types by _import_plugin_decorators().
707
+ """
708
+ ...
709
+
710
+ @typing.overload
711
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
712
+ ...
713
+
714
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = 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 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]]]:
723
+ """
724
+ Creates a human-readable report, a Metaflow Card, after this step completes.
670
725
 
671
- Information in this decorator will augment any
672
- attributes set in the `@pyi_base` flow-level decorator. Hence,
673
- you can use `@pypi_base` to set packages required by all
674
- steps and use `@pypi` to specify step-specific overrides.
726
+ Note that you may add multiple `@card` decorators in a step with different parameters.
675
727
 
676
728
 
677
729
  Parameters
678
730
  ----------
679
- packages : Dict[str, str], default: {}
680
- Packages to use for this step. The key is the name of the package
681
- and the value is the version to use.
682
- python : str, optional, default: None
683
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
684
- that the version used will correspond to the version of the Python interpreter used to start the run.
731
+ type : str, default 'default'
732
+ Card type.
733
+ id : str, optional, default None
734
+ If multiple cards are present, use this id to identify this card.
735
+ options : Dict[str, Any], default {}
736
+ Options passed to the card. The contents depend on the card type.
737
+ timeout : int, default 45
738
+ Interrupt reporting if it takes more than this many seconds.
685
739
  """
686
740
  ...
687
741
 
688
742
  @typing.overload
689
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
743
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
690
744
  ...
691
745
 
692
746
  @typing.overload
693
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
747
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
694
748
  ...
695
749
 
696
- 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):
750
+ 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):
697
751
  """
698
- Specifies the PyPI packages for the step.
752
+ Creates a human-readable report, a Metaflow Card, after this step completes.
699
753
 
700
- Information in this decorator will augment any
701
- attributes set in the `@pyi_base` flow-level decorator. Hence,
702
- you can use `@pypi_base` to set packages required by all
703
- steps and use `@pypi` to specify step-specific overrides.
754
+ Note that you may add multiple `@card` decorators in a step with different parameters.
704
755
 
705
756
 
706
757
  Parameters
707
758
  ----------
708
- packages : Dict[str, str], default: {}
709
- Packages to use for this step. The key is the name of the package
710
- and the value is the version to use.
711
- python : str, optional, default: None
712
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
713
- that the version used will correspond to the version of the Python interpreter used to start the run.
759
+ type : str, default 'default'
760
+ Card type.
761
+ id : str, optional, default None
762
+ If multiple cards are present, use this id to identify this card.
763
+ options : Dict[str, Any], default {}
764
+ Options passed to the card. The contents depend on the card type.
765
+ timeout : int, default 45
766
+ Interrupt reporting if it takes more than this many seconds.
714
767
  """
715
768
  ...
716
769
 
717
770
  @typing.overload
718
- 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]]]:
771
+ def catch(*, var: typing.Optional[str] = None, print_exception: bool = True) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
719
772
  """
720
- Enables checkpointing for a step.
773
+ Specifies that the step will success under all circumstances.
721
774
 
722
- > Examples
775
+ The decorator will create an optional artifact, specified by `var`, which
776
+ contains the exception raised. You can use it to detect the presence
777
+ of errors, indicating that all happy-path artifacts produced by the step
778
+ are missing.
723
779
 
724
- - Saving Checkpoints
725
780
 
726
- ```python
727
- @checkpoint
728
- @step
729
- def train(self):
730
- model = create_model(self.parameters, checkpoint_path = None)
731
- for i in range(self.epochs):
732
- # some training logic
733
- loss = model.train(self.dataset)
734
- if i % 10 == 0:
735
- model.save(
736
- current.checkpoint.directory,
781
+ Parameters
782
+ ----------
783
+ var : str, optional, default None
784
+ Name of the artifact in which to store the caught exception.
785
+ If not specified, the exception is not stored.
786
+ print_exception : bool, default True
787
+ Determines whether or not the exception is printed to
788
+ stdout when caught.
789
+ """
790
+ ...
791
+
792
+ @typing.overload
793
+ def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
794
+ ...
795
+
796
+ @typing.overload
797
+ def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
798
+ ...
799
+
800
+ def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, var: typing.Optional[str] = None, print_exception: bool = True):
801
+ """
802
+ Specifies that the step will success under all circumstances.
803
+
804
+ The decorator will create an optional artifact, specified by `var`, which
805
+ contains the exception raised. You can use it to detect the presence
806
+ of errors, indicating that all happy-path artifacts produced by the step
807
+ are missing.
808
+
809
+
810
+ Parameters
811
+ ----------
812
+ var : str, optional, default None
813
+ Name of the artifact in which to store the caught exception.
814
+ If not specified, the exception is not stored.
815
+ print_exception : bool, default True
816
+ Determines whether or not the exception is printed to
817
+ stdout when caught.
818
+ """
819
+ ...
820
+
821
+ @typing.overload
822
+ 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]]]:
823
+ """
824
+ Enables loading / saving of models within a step.
825
+
826
+ > Examples
827
+ - Saving Models
828
+ ```python
829
+ @model
830
+ @step
831
+ def train(self):
832
+ # current.model.save returns a dictionary reference to the model saved
833
+ self.my_model = current.model.save(
834
+ path_to_my_model,
835
+ label="my_model",
836
+ metadata={
837
+ "epochs": 10,
838
+ "batch-size": 32,
839
+ "learning-rate": 0.001,
840
+ }
841
+ )
842
+ self.next(self.test)
843
+
844
+ @model(load="my_model")
845
+ @step
846
+ def test(self):
847
+ # `current.model.loaded` returns a dictionary of the loaded models
848
+ # where the key is the name of the artifact and the value is the path to the model
849
+ print(os.listdir(current.model.loaded["my_model"]))
850
+ self.next(self.end)
851
+ ```
852
+
853
+ - Loading models
854
+ ```python
855
+ @step
856
+ def train(self):
857
+ # current.model.load returns the path to the model loaded
858
+ checkpoint_path = current.model.load(
859
+ self.checkpoint_key,
860
+ )
861
+ model_path = current.model.load(
862
+ self.model,
863
+ )
864
+ self.next(self.test)
865
+ ```
866
+
867
+
868
+ Parameters
869
+ ----------
870
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
871
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
872
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by `current.checkpoint` / `current.model` / `current.huggingface_hub`.
873
+ If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
874
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
875
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
876
+
877
+ temp_dir_root : str, default: None
878
+ The root directory under which `current.model.loaded` will store loaded models
879
+ """
880
+ ...
881
+
882
+ @typing.overload
883
+ def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
884
+ ...
885
+
886
+ @typing.overload
887
+ def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
888
+ ...
889
+
890
+ def model(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, load: typing.Union[typing.List[str], str, typing.List[typing.Tuple[str, typing.Optional[str]]]] = None, temp_dir_root: str = None):
891
+ """
892
+ Enables loading / saving of models within a step.
893
+
894
+ > Examples
895
+ - Saving Models
896
+ ```python
897
+ @model
898
+ @step
899
+ def train(self):
900
+ # current.model.save returns a dictionary reference to the model saved
901
+ self.my_model = current.model.save(
902
+ path_to_my_model,
903
+ label="my_model",
904
+ metadata={
905
+ "epochs": 10,
906
+ "batch-size": 32,
907
+ "learning-rate": 0.001,
908
+ }
909
+ )
910
+ self.next(self.test)
911
+
912
+ @model(load="my_model")
913
+ @step
914
+ def test(self):
915
+ # `current.model.loaded` returns a dictionary of the loaded models
916
+ # where the key is the name of the artifact and the value is the path to the model
917
+ print(os.listdir(current.model.loaded["my_model"]))
918
+ self.next(self.end)
919
+ ```
920
+
921
+ - Loading models
922
+ ```python
923
+ @step
924
+ def train(self):
925
+ # current.model.load returns the path to the model loaded
926
+ checkpoint_path = current.model.load(
927
+ self.checkpoint_key,
928
+ )
929
+ model_path = current.model.load(
930
+ self.model,
931
+ )
932
+ self.next(self.test)
933
+ ```
934
+
935
+
936
+ Parameters
937
+ ----------
938
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
939
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
940
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by `current.checkpoint` / `current.model` / `current.huggingface_hub`.
941
+ If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
942
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
943
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
944
+
945
+ temp_dir_root : str, default: None
946
+ The root directory under which `current.model.loaded` will store loaded models
947
+ """
948
+ ...
949
+
950
+ @typing.overload
951
+ 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]]]:
952
+ """
953
+ Enables checkpointing for a step.
954
+
955
+ > Examples
956
+
957
+ - Saving Checkpoints
958
+
959
+ ```python
960
+ @checkpoint
961
+ @step
962
+ def train(self):
963
+ model = create_model(self.parameters, checkpoint_path = None)
964
+ for i in range(self.epochs):
965
+ # some training logic
966
+ loss = model.train(self.dataset)
967
+ if i % 10 == 0:
968
+ model.save(
969
+ current.checkpoint.directory,
737
970
  )
738
971
  # saves the contents of the `current.checkpoint.directory` as a checkpoint
739
972
  # and returns a reference dictionary to the checkpoint saved in the datastore
@@ -862,388 +1095,80 @@ def checkpoint(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None
862
1095
  ...
863
1096
 
864
1097
  @typing.overload
865
- def timeout(*, seconds: int = 0, minutes: int = 0, hours: int = 0) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
866
- """
867
- Specifies a timeout for your step.
868
-
869
- This decorator is useful if this step may hang indefinitely.
870
-
871
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
872
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
873
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
874
-
875
- Note that all the values specified in parameters are added together so if you specify
876
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
877
-
878
-
879
- Parameters
880
- ----------
881
- seconds : int, default 0
882
- Number of seconds to wait prior to timing out.
883
- minutes : int, default 0
884
- Number of minutes to wait prior to timing out.
885
- hours : int, default 0
886
- Number of hours to wait prior to timing out.
887
- """
888
- ...
889
-
890
- @typing.overload
891
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
892
- ...
893
-
894
- @typing.overload
895
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
896
- ...
897
-
898
- def timeout(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, seconds: int = 0, minutes: int = 0, hours: int = 0):
899
- """
900
- Specifies a timeout for your step.
901
-
902
- This decorator is useful if this step may hang indefinitely.
903
-
904
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
905
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
906
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
907
-
908
- Note that all the values specified in parameters are added together so if you specify
909
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
910
-
911
-
912
- Parameters
913
- ----------
914
- seconds : int, default 0
915
- Number of seconds to wait prior to timing out.
916
- minutes : int, default 0
917
- Number of minutes to wait prior to timing out.
918
- hours : int, default 0
919
- Number of hours to wait prior to timing out.
920
- """
921
- ...
922
-
923
- @typing.overload
924
- def test_append_card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
925
- """
926
- A simple decorator that demonstrates using CardDecoratorInjector
927
- to inject a card and render simple markdown content.
928
- """
929
- ...
930
-
931
- @typing.overload
932
- def test_append_card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
933
- ...
934
-
935
- def test_append_card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
936
- """
937
- A simple decorator that demonstrates using CardDecoratorInjector
938
- to inject a card and render simple markdown content.
939
- """
940
- ...
941
-
942
- @typing.overload
943
- 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]]]:
944
- """
945
- Specifies the resources needed when executing this step.
946
-
947
- Use `@resources` to specify the resource requirements
948
- independently of the specific compute layer (`@batch`, `@kubernetes`).
949
-
950
- You can choose the compute layer on the command line by executing e.g.
951
- ```
952
- python myflow.py run --with batch
953
- ```
954
- or
955
- ```
956
- python myflow.py run --with kubernetes
957
- ```
958
- which executes the flow on the desired system using the
959
- requirements specified in `@resources`.
960
-
961
-
962
- Parameters
963
- ----------
964
- cpu : int, default 1
965
- Number of CPUs required for this step.
966
- gpu : int, optional, default None
967
- Number of GPUs required for this step.
968
- disk : int, optional, default None
969
- Disk size (in MB) required for this step. Only applies on Kubernetes.
970
- memory : int, default 4096
971
- Memory size (in MB) required for this step.
972
- shared_memory : int, optional, default None
973
- The value for the size (in MiB) of the /dev/shm volume for this step.
974
- This parameter maps to the `--shm-size` option in Docker.
975
- """
976
- ...
977
-
978
- @typing.overload
979
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
980
- ...
981
-
982
- @typing.overload
983
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
984
- ...
985
-
986
- 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):
987
- """
988
- Specifies the resources needed when executing this step.
989
-
990
- Use `@resources` to specify the resource requirements
991
- independently of the specific compute layer (`@batch`, `@kubernetes`).
992
-
993
- You can choose the compute layer on the command line by executing e.g.
994
- ```
995
- python myflow.py run --with batch
996
- ```
997
- or
998
- ```
999
- python myflow.py run --with kubernetes
1000
- ```
1001
- which executes the flow on the desired system using the
1002
- requirements specified in `@resources`.
1003
-
1004
-
1005
- Parameters
1006
- ----------
1007
- cpu : int, default 1
1008
- Number of CPUs required for this step.
1009
- gpu : int, optional, default None
1010
- Number of GPUs required for this step.
1011
- disk : int, optional, default None
1012
- Disk size (in MB) required for this step. Only applies on Kubernetes.
1013
- memory : int, default 4096
1014
- Memory size (in MB) required for this step.
1015
- shared_memory : int, optional, default None
1016
- The value for the size (in MiB) of the /dev/shm volume for this step.
1017
- This parameter maps to the `--shm-size` option in Docker.
1018
- """
1019
- ...
1020
-
1021
- @typing.overload
1022
- def catch(*, var: typing.Optional[str] = None, print_exception: bool = True) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
1023
- """
1024
- Specifies that the step will success under all circumstances.
1025
-
1026
- The decorator will create an optional artifact, specified by `var`, which
1027
- contains the exception raised. You can use it to detect the presence
1028
- of errors, indicating that all happy-path artifacts produced by the step
1029
- are missing.
1030
-
1031
-
1032
- Parameters
1033
- ----------
1034
- var : str, optional, default None
1035
- Name of the artifact in which to store the caught exception.
1036
- If not specified, the exception is not stored.
1037
- print_exception : bool, default True
1038
- Determines whether or not the exception is printed to
1039
- stdout when caught.
1040
- """
1041
- ...
1042
-
1043
- @typing.overload
1044
- def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1045
- ...
1046
-
1047
- @typing.overload
1048
- def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1049
- ...
1050
-
1051
- def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, var: typing.Optional[str] = None, print_exception: bool = True):
1052
- """
1053
- Specifies that the step will success under all circumstances.
1054
-
1055
- The decorator will create an optional artifact, specified by `var`, which
1056
- contains the exception raised. You can use it to detect the presence
1057
- of errors, indicating that all happy-path artifacts produced by the step
1058
- are missing.
1059
-
1060
-
1061
- Parameters
1062
- ----------
1063
- var : str, optional, default None
1064
- Name of the artifact in which to store the caught exception.
1065
- If not specified, the exception is not stored.
1066
- print_exception : bool, default True
1067
- Determines whether or not the exception is printed to
1068
- stdout when caught.
1069
- """
1070
- ...
1071
-
1072
- 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]]]:
1098
+ 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]]]:
1073
1099
  """
1074
- Specifies that this step should execute on DGX cloud.
1075
-
1100
+ Specifies the Conda environment for the step.
1076
1101
 
1077
- Parameters
1078
- ----------
1079
- gpu : int
1080
- Number of GPUs to use.
1081
- gpu_type : str
1082
- Type of Nvidia GPU to use.
1083
- queue_timeout : int
1084
- Time to keep the job in NVCF's queue.
1085
- """
1086
- ...
1087
-
1088
- @typing.overload
1089
- 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]]]:
1090
- """
1091
- Specifies environment variables to be set prior to the execution of a step.
1102
+ Information in this decorator will augment any
1103
+ attributes set in the `@conda_base` flow-level decorator. Hence,
1104
+ you can use `@conda_base` to set packages required by all
1105
+ steps and use `@conda` to specify step-specific overrides.
1092
1106
 
1093
1107
 
1094
1108
  Parameters
1095
1109
  ----------
1096
- vars : Dict[str, str], default {}
1097
- Dictionary of environment variables to set.
1110
+ packages : Dict[str, str], default {}
1111
+ Packages to use for this step. The key is the name of the package
1112
+ and the value is the version to use.
1113
+ libraries : Dict[str, str], default {}
1114
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1115
+ python : str, optional, default None
1116
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1117
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1118
+ disabled : bool, default False
1119
+ If set to True, disables @conda.
1098
1120
  """
1099
1121
  ...
1100
1122
 
1101
1123
  @typing.overload
1102
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1124
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1103
1125
  ...
1104
1126
 
1105
1127
  @typing.overload
1106
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1107
- ...
1108
-
1109
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
1110
- """
1111
- Specifies environment variables to be set prior to the execution of a step.
1112
-
1113
-
1114
- Parameters
1115
- ----------
1116
- vars : Dict[str, str], default {}
1117
- Dictionary of environment variables to set.
1118
- """
1128
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1119
1129
  ...
1120
1130
 
1121
- @typing.overload
1122
- 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]]]:
1131
+ 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):
1123
1132
  """
1124
- Enables loading / saving of models within a step.
1125
-
1126
- > Examples
1127
- - Saving Models
1128
- ```python
1129
- @model
1130
- @step
1131
- def train(self):
1132
- # current.model.save returns a dictionary reference to the model saved
1133
- self.my_model = current.model.save(
1134
- path_to_my_model,
1135
- label="my_model",
1136
- metadata={
1137
- "epochs": 10,
1138
- "batch-size": 32,
1139
- "learning-rate": 0.001,
1140
- }
1141
- )
1142
- self.next(self.test)
1143
-
1144
- @model(load="my_model")
1145
- @step
1146
- def test(self):
1147
- # `current.model.loaded` returns a dictionary of the loaded models
1148
- # where the key is the name of the artifact and the value is the path to the model
1149
- print(os.listdir(current.model.loaded["my_model"]))
1150
- self.next(self.end)
1151
- ```
1133
+ Specifies the Conda environment for the step.
1152
1134
 
1153
- - Loading models
1154
- ```python
1155
- @step
1156
- def train(self):
1157
- # current.model.load returns the path to the model loaded
1158
- checkpoint_path = current.model.load(
1159
- self.checkpoint_key,
1160
- )
1161
- model_path = current.model.load(
1162
- self.model,
1163
- )
1164
- self.next(self.test)
1165
- ```
1135
+ Information in this decorator will augment any
1136
+ attributes set in the `@conda_base` flow-level decorator. Hence,
1137
+ you can use `@conda_base` to set packages required by all
1138
+ steps and use `@conda` to specify step-specific overrides.
1166
1139
 
1167
1140
 
1168
1141
  Parameters
1169
1142
  ----------
1170
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
1171
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
1172
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by `current.checkpoint` / `current.model` / `current.huggingface_hub`.
1173
- If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
1174
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
1175
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
1176
-
1177
- temp_dir_root : str, default: None
1178
- The root directory under which `current.model.loaded` will store loaded models
1143
+ packages : Dict[str, str], default {}
1144
+ Packages to use for this step. The key is the name of the package
1145
+ and the value is the version to use.
1146
+ libraries : Dict[str, str], default {}
1147
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1148
+ python : str, optional, default None
1149
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1150
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1151
+ disabled : bool, default False
1152
+ If set to True, disables @conda.
1179
1153
  """
1180
1154
  ...
1181
1155
 
1182
1156
  @typing.overload
1183
- def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1157
+ def app_deploy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], 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
+ """
1184
1162
  ...
1185
1163
 
1186
1164
  @typing.overload
1187
- def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1165
+ def app_deploy(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1188
1166
  ...
1189
1167
 
1190
- def model(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, load: typing.Union[typing.List[str], str, typing.List[typing.Tuple[str, typing.Optional[str]]]] = None, temp_dir_root: str = None):
1168
+ def app_deploy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1191
1169
  """
1192
- Enables loading / saving of models within a step.
1193
-
1194
- > Examples
1195
- - Saving Models
1196
- ```python
1197
- @model
1198
- @step
1199
- def train(self):
1200
- # current.model.save returns a dictionary reference to the model saved
1201
- self.my_model = current.model.save(
1202
- path_to_my_model,
1203
- label="my_model",
1204
- metadata={
1205
- "epochs": 10,
1206
- "batch-size": 32,
1207
- "learning-rate": 0.001,
1208
- }
1209
- )
1210
- self.next(self.test)
1211
-
1212
- @model(load="my_model")
1213
- @step
1214
- def test(self):
1215
- # `current.model.loaded` returns a dictionary of the loaded models
1216
- # where the key is the name of the artifact and the value is the path to the model
1217
- print(os.listdir(current.model.loaded["my_model"]))
1218
- self.next(self.end)
1219
- ```
1220
-
1221
- - Loading models
1222
- ```python
1223
- @step
1224
- def train(self):
1225
- # current.model.load returns the path to the model loaded
1226
- checkpoint_path = current.model.load(
1227
- self.checkpoint_key,
1228
- )
1229
- model_path = current.model.load(
1230
- self.model,
1231
- )
1232
- self.next(self.test)
1233
- ```
1234
-
1235
-
1236
- Parameters
1237
- ----------
1238
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
1239
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
1240
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by `current.checkpoint` / `current.model` / `current.huggingface_hub`.
1241
- If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
1242
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
1243
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
1244
-
1245
- temp_dir_root : str, default: None
1246
- The root directory under which `current.model.loaded` will store loaded models
1170
+ Decorator prototype for all step decorators. This function gets specialized
1171
+ and imported for all decorators types by _import_plugin_decorators().
1247
1172
  """
1248
1173
  ...
1249
1174
 
@@ -1337,248 +1262,136 @@ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: ty
1337
1262
  ...
1338
1263
 
1339
1264
  @typing.overload
1340
- 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]]]:
1265
+ 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]]]:
1341
1266
  """
1342
- Specifies the Conda environment for the step.
1267
+ Specifies the PyPI packages for the step.
1343
1268
 
1344
1269
  Information in this decorator will augment any
1345
- attributes set in the `@conda_base` flow-level decorator. Hence,
1346
- you can use `@conda_base` to set packages required by all
1347
- steps and use `@conda` to specify step-specific overrides.
1270
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
1271
+ you can use `@pypi_base` to set packages required by all
1272
+ steps and use `@pypi` to specify step-specific overrides.
1348
1273
 
1349
1274
 
1350
1275
  Parameters
1351
1276
  ----------
1352
- packages : Dict[str, str], default {}
1277
+ packages : Dict[str, str], default: {}
1353
1278
  Packages to use for this step. The key is the name of the package
1354
1279
  and the value is the version to use.
1355
- libraries : Dict[str, str], default {}
1356
- Supported for backward compatibility. When used with packages, packages will take precedence.
1357
- python : str, optional, default None
1280
+ python : str, optional, default: None
1358
1281
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
1359
1282
  that the version used will correspond to the version of the Python interpreter used to start the run.
1360
- disabled : bool, default False
1361
- If set to True, disables @conda.
1362
1283
  """
1363
1284
  ...
1364
1285
 
1365
1286
  @typing.overload
1366
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1287
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1367
1288
  ...
1368
1289
 
1369
1290
  @typing.overload
1370
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1291
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1371
1292
  ...
1372
1293
 
1373
- 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):
1294
+ 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):
1374
1295
  """
1375
- Specifies the Conda environment for the step.
1296
+ Specifies the PyPI packages for the step.
1376
1297
 
1377
1298
  Information in this decorator will augment any
1378
- attributes set in the `@conda_base` flow-level decorator. Hence,
1379
- you can use `@conda_base` to set packages required by all
1380
- steps and use `@conda` to specify step-specific overrides.
1299
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
1300
+ you can use `@pypi_base` to set packages required by all
1301
+ steps and use `@pypi` to specify step-specific overrides.
1381
1302
 
1382
1303
 
1383
1304
  Parameters
1384
1305
  ----------
1385
- packages : Dict[str, str], default {}
1306
+ packages : Dict[str, str], default: {}
1386
1307
  Packages to use for this step. The key is the name of the package
1387
1308
  and the value is the version to use.
1388
- libraries : Dict[str, str], default {}
1389
- Supported for backward compatibility. When used with packages, packages will take precedence.
1390
- python : str, optional, default None
1391
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1392
- that the version used will correspond to the version of the Python interpreter used to start the run.
1393
- disabled : bool, default False
1394
- If set to True, disables @conda.
1395
- """
1396
- ...
1397
-
1398
- @typing.overload
1399
- def conda_base(*, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1400
- """
1401
- Specifies the Conda environment for all steps of the flow.
1402
-
1403
- Use `@conda_base` to set common libraries required by all
1404
- steps and use `@conda` to specify step-specific additions.
1405
-
1406
-
1407
- Parameters
1408
- ----------
1409
- packages : Dict[str, str], default {}
1410
- Packages to use for this flow. The key is the name of the package
1411
- and the value is the version to use.
1412
- libraries : Dict[str, str], default {}
1413
- Supported for backward compatibility. When used with packages, packages will take precedence.
1414
- python : str, optional, default None
1309
+ python : str, optional, default: None
1415
1310
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
1416
1311
  that the version used will correspond to the version of the Python interpreter used to start the run.
1417
- disabled : bool, default False
1418
- If set to True, disables Conda.
1419
1312
  """
1420
1313
  ...
1421
1314
 
1422
1315
  @typing.overload
1423
- def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1424
- ...
1425
-
1426
- def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
1316
+ 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]]]:
1427
1317
  """
1428
- Specifies the Conda environment for all steps of the flow.
1429
-
1430
- Use `@conda_base` to set common libraries required by all
1431
- steps and use `@conda` to specify step-specific additions.
1318
+ Specifies environment variables to be set prior to the execution of a step.
1432
1319
 
1433
1320
 
1434
1321
  Parameters
1435
1322
  ----------
1436
- packages : Dict[str, str], default {}
1437
- Packages to use for this flow. The key is the name of the package
1438
- and the value is the version to use.
1439
- libraries : Dict[str, str], default {}
1440
- Supported for backward compatibility. When used with packages, packages will take precedence.
1441
- python : str, optional, default None
1442
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1443
- that the version used will correspond to the version of the Python interpreter used to start the run.
1444
- disabled : bool, default False
1445
- If set to True, disables Conda.
1323
+ vars : Dict[str, str], default {}
1324
+ Dictionary of environment variables to set.
1446
1325
  """
1447
1326
  ...
1448
1327
 
1449
1328
  @typing.overload
1450
- 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]]:
1451
- """
1452
- Specifies the event(s) that this flow depends on.
1453
-
1454
- ```
1455
- @trigger(event='foo')
1456
- ```
1457
- or
1458
- ```
1459
- @trigger(events=['foo', 'bar'])
1460
- ```
1461
-
1462
- Additionally, you can specify the parameter mappings
1463
- to map event payload to Metaflow parameters for the flow.
1464
- ```
1465
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1466
- ```
1467
- or
1468
- ```
1469
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1470
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1471
- ```
1472
-
1473
- 'parameters' can also be a list of strings and tuples like so:
1474
- ```
1475
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1476
- ```
1477
- This is equivalent to:
1478
- ```
1479
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1480
- ```
1481
-
1482
-
1483
- Parameters
1484
- ----------
1485
- event : Union[str, Dict[str, Any]], optional, default None
1486
- Event dependency for this flow.
1487
- events : List[Union[str, Dict[str, Any]]], default []
1488
- Events dependency for this flow.
1489
- options : Dict[str, Any], default {}
1490
- Backend-specific configuration for tuning eventing behavior.
1491
- """
1329
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1492
1330
  ...
1493
1331
 
1494
1332
  @typing.overload
1495
- def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1333
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1496
1334
  ...
1497
1335
 
1498
- def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: typing.Union[str, typing.Dict[str, typing.Any], None] = None, events: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], options: typing.Dict[str, typing.Any] = {}):
1336
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
1499
1337
  """
1500
- Specifies the event(s) that this flow depends on.
1501
-
1502
- ```
1503
- @trigger(event='foo')
1504
- ```
1505
- or
1506
- ```
1507
- @trigger(events=['foo', 'bar'])
1508
- ```
1509
-
1510
- Additionally, you can specify the parameter mappings
1511
- to map event payload to Metaflow parameters for the flow.
1512
- ```
1513
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1514
- ```
1515
- or
1516
- ```
1517
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1518
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1519
- ```
1520
-
1521
- 'parameters' can also be a list of strings and tuples like so:
1522
- ```
1523
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1524
- ```
1525
- This is equivalent to:
1526
- ```
1527
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1528
- ```
1338
+ Specifies environment variables to be set prior to the execution of a step.
1529
1339
 
1530
1340
 
1531
1341
  Parameters
1532
1342
  ----------
1533
- event : Union[str, Dict[str, Any]], optional, default None
1534
- Event dependency for this flow.
1535
- events : List[Union[str, Dict[str, Any]]], default []
1536
- Events dependency for this flow.
1537
- options : Dict[str, Any], default {}
1538
- Backend-specific configuration for tuning eventing behavior.
1343
+ vars : Dict[str, str], default {}
1344
+ Dictionary of environment variables to set.
1539
1345
  """
1540
1346
  ...
1541
1347
 
1542
- 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]]:
1348
+ 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]]]:
1543
1349
  """
1544
- 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.
1545
- 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.
1350
+ This decorator is used to run vllm APIs as Metaflow task sidecars.
1351
+
1352
+ User code call
1353
+ --------------
1354
+ @vllm(
1355
+ model="...",
1356
+ ...
1357
+ )
1358
+
1359
+ Valid backend options
1360
+ ---------------------
1361
+ - 'local': Run as a separate process on the local task machine.
1362
+
1363
+ Valid model options
1364
+ -------------------
1365
+ Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
1366
+
1367
+ NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
1368
+ If you need multiple models, you must create multiple @vllm decorators.
1546
1369
 
1547
1370
 
1548
1371
  Parameters
1549
1372
  ----------
1550
- timeout : int
1551
- Time, in seconds before the task times out and fails. (Default: 3600)
1552
- poke_interval : int
1553
- Time in seconds that the job should wait in between each try. (Default: 60)
1554
- mode : str
1555
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1556
- exponential_backoff : bool
1557
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1558
- pool : str
1559
- the slot pool this task should run in,
1560
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1561
- soft_fail : bool
1562
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1563
- name : str
1564
- Name of the sensor on Airflow
1565
- description : str
1566
- Description of sensor in the Airflow UI
1567
- external_dag_id : str
1568
- The dag_id that contains the task you want to wait for.
1569
- external_task_ids : List[str]
1570
- The list of task_ids that you want to wait for.
1571
- If None (default value) the sensor waits for the DAG. (Default: None)
1572
- allowed_states : List[str]
1573
- Iterable of allowed states, (Default: ['success'])
1574
- failed_states : List[str]
1575
- Iterable of failed or dis-allowed states. (Default: None)
1576
- execution_delta : datetime.timedelta
1577
- time difference with the previous execution to look at,
1578
- the default is the same logical date as the current task or DAG. (Default: None)
1579
- check_existence: bool
1580
- Set to True to check if the external task exists or check if
1581
- the DAG to wait for exists. (Default: True)
1373
+ model: str
1374
+ HuggingFace model identifier to be served by vLLM.
1375
+ backend: str
1376
+ Determines where and how to run the vLLM process.
1377
+ openai_api_server: bool
1378
+ Whether to use OpenAI-compatible API server mode (subprocess) instead of native engine.
1379
+ Default is False (uses native engine).
1380
+ Set to True for backward compatibility with existing code.
1381
+ debug: bool
1382
+ Whether to turn on verbose debugging logs.
1383
+ card_refresh_interval: int
1384
+ Interval in seconds for refreshing the vLLM status card.
1385
+ Only used when openai_api_server=True.
1386
+ max_retries: int
1387
+ Maximum number of retries checking for vLLM server startup.
1388
+ Only used when openai_api_server=True.
1389
+ retry_alert_frequency: int
1390
+ Frequency of alert logs for vLLM server startup retries.
1391
+ Only used when openai_api_server=True.
1392
+ engine_args : dict
1393
+ Additional keyword arguments to pass to the vLLM engine.
1394
+ For example, `tensor_parallel_size=2`.
1582
1395
  """
1583
1396
  ...
1584
1397
 
@@ -1967,5 +1780,192 @@ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packag
1967
1780
  """
1968
1781
  ...
1969
1782
 
1783
+ 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]]:
1784
+ """
1785
+ 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.
1786
+ 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.
1787
+
1788
+
1789
+ Parameters
1790
+ ----------
1791
+ timeout : int
1792
+ Time, in seconds before the task times out and fails. (Default: 3600)
1793
+ poke_interval : int
1794
+ Time in seconds that the job should wait in between each try. (Default: 60)
1795
+ mode : str
1796
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1797
+ exponential_backoff : bool
1798
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1799
+ pool : str
1800
+ the slot pool this task should run in,
1801
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1802
+ soft_fail : bool
1803
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1804
+ name : str
1805
+ Name of the sensor on Airflow
1806
+ description : str
1807
+ Description of sensor in the Airflow UI
1808
+ external_dag_id : str
1809
+ The dag_id that contains the task you want to wait for.
1810
+ external_task_ids : List[str]
1811
+ The list of task_ids that you want to wait for.
1812
+ If None (default value) the sensor waits for the DAG. (Default: None)
1813
+ allowed_states : List[str]
1814
+ Iterable of allowed states, (Default: ['success'])
1815
+ failed_states : List[str]
1816
+ Iterable of failed or dis-allowed states. (Default: None)
1817
+ execution_delta : datetime.timedelta
1818
+ time difference with the previous execution to look at,
1819
+ the default is the same logical date as the current task or DAG. (Default: None)
1820
+ check_existence: bool
1821
+ Set to True to check if the external task exists or check if
1822
+ the DAG to wait for exists. (Default: True)
1823
+ """
1824
+ ...
1825
+
1826
+ @typing.overload
1827
+ def conda_base(*, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1828
+ """
1829
+ Specifies the Conda environment for all steps of the flow.
1830
+
1831
+ Use `@conda_base` to set common libraries required by all
1832
+ steps and use `@conda` to specify step-specific additions.
1833
+
1834
+
1835
+ Parameters
1836
+ ----------
1837
+ packages : Dict[str, str], default {}
1838
+ Packages to use for this flow. The key is the name of the package
1839
+ and the value is the version to use.
1840
+ libraries : Dict[str, str], default {}
1841
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1842
+ python : str, optional, default None
1843
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1844
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1845
+ disabled : bool, default False
1846
+ If set to True, disables Conda.
1847
+ """
1848
+ ...
1849
+
1850
+ @typing.overload
1851
+ def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1852
+ ...
1853
+
1854
+ def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
1855
+ """
1856
+ Specifies the Conda environment for all steps of the flow.
1857
+
1858
+ Use `@conda_base` to set common libraries required by all
1859
+ steps and use `@conda` to specify step-specific additions.
1860
+
1861
+
1862
+ Parameters
1863
+ ----------
1864
+ packages : Dict[str, str], default {}
1865
+ Packages to use for this flow. The key is the name of the package
1866
+ and the value is the version to use.
1867
+ libraries : Dict[str, str], default {}
1868
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1869
+ python : str, optional, default None
1870
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1871
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1872
+ disabled : bool, default False
1873
+ If set to True, disables Conda.
1874
+ """
1875
+ ...
1876
+
1877
+ @typing.overload
1878
+ 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]]:
1879
+ """
1880
+ Specifies the event(s) that this flow depends on.
1881
+
1882
+ ```
1883
+ @trigger(event='foo')
1884
+ ```
1885
+ or
1886
+ ```
1887
+ @trigger(events=['foo', 'bar'])
1888
+ ```
1889
+
1890
+ Additionally, you can specify the parameter mappings
1891
+ to map event payload to Metaflow parameters for the flow.
1892
+ ```
1893
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1894
+ ```
1895
+ or
1896
+ ```
1897
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1898
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1899
+ ```
1900
+
1901
+ 'parameters' can also be a list of strings and tuples like so:
1902
+ ```
1903
+ @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1904
+ ```
1905
+ This is equivalent to:
1906
+ ```
1907
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1908
+ ```
1909
+
1910
+
1911
+ Parameters
1912
+ ----------
1913
+ event : Union[str, Dict[str, Any]], optional, default None
1914
+ Event dependency for this flow.
1915
+ events : List[Union[str, Dict[str, Any]]], default []
1916
+ Events dependency for this flow.
1917
+ options : Dict[str, Any], default {}
1918
+ Backend-specific configuration for tuning eventing behavior.
1919
+ """
1920
+ ...
1921
+
1922
+ @typing.overload
1923
+ def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1924
+ ...
1925
+
1926
+ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: typing.Union[str, typing.Dict[str, typing.Any], None] = None, events: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], options: typing.Dict[str, typing.Any] = {}):
1927
+ """
1928
+ Specifies the event(s) that this flow depends on.
1929
+
1930
+ ```
1931
+ @trigger(event='foo')
1932
+ ```
1933
+ or
1934
+ ```
1935
+ @trigger(events=['foo', 'bar'])
1936
+ ```
1937
+
1938
+ Additionally, you can specify the parameter mappings
1939
+ to map event payload to Metaflow parameters for the flow.
1940
+ ```
1941
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1942
+ ```
1943
+ or
1944
+ ```
1945
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1946
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1947
+ ```
1948
+
1949
+ 'parameters' can also be a list of strings and tuples like so:
1950
+ ```
1951
+ @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1952
+ ```
1953
+ This is equivalent to:
1954
+ ```
1955
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1956
+ ```
1957
+
1958
+
1959
+ Parameters
1960
+ ----------
1961
+ event : Union[str, Dict[str, Any]], optional, default None
1962
+ Event dependency for this flow.
1963
+ events : List[Union[str, Dict[str, Any]]], default []
1964
+ Events dependency for this flow.
1965
+ options : Dict[str, Any], default {}
1966
+ Backend-specific configuration for tuning eventing behavior.
1967
+ """
1968
+ ...
1969
+
1970
1970
  pkg_name: str
1971
1971