ob-metaflow-extensions 1.1.171rc1__py2.py3-none-any.whl → 1.4.35__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-extensions might be problematic. Click here for more details.

Files changed (64) hide show
  1. metaflow_extensions/outerbounds/plugins/__init__.py +6 -3
  2. metaflow_extensions/outerbounds/plugins/apps/app_cli.py +0 -29
  3. metaflow_extensions/outerbounds/plugins/apps/app_deploy_decorator.py +146 -0
  4. metaflow_extensions/outerbounds/plugins/apps/core/__init__.py +10 -0
  5. metaflow_extensions/outerbounds/plugins/apps/core/_state_machine.py +506 -0
  6. metaflow_extensions/outerbounds/plugins/apps/core/_vendor/__init__.py +0 -0
  7. metaflow_extensions/outerbounds/plugins/apps/core/_vendor/spinner/__init__.py +4 -0
  8. metaflow_extensions/outerbounds/plugins/apps/core/_vendor/spinner/spinners.py +478 -0
  9. metaflow_extensions/outerbounds/plugins/apps/core/app_cli.py +1200 -0
  10. metaflow_extensions/outerbounds/plugins/apps/core/app_config.py +146 -0
  11. metaflow_extensions/outerbounds/plugins/apps/core/artifacts.py +0 -0
  12. metaflow_extensions/outerbounds/plugins/apps/core/capsule.py +958 -0
  13. metaflow_extensions/outerbounds/plugins/apps/core/click_importer.py +24 -0
  14. metaflow_extensions/outerbounds/plugins/apps/core/code_package/__init__.py +3 -0
  15. metaflow_extensions/outerbounds/plugins/apps/core/code_package/code_packager.py +618 -0
  16. metaflow_extensions/outerbounds/plugins/apps/core/code_package/examples.py +125 -0
  17. metaflow_extensions/outerbounds/plugins/apps/core/config/__init__.py +12 -0
  18. metaflow_extensions/outerbounds/plugins/apps/core/config/cli_generator.py +161 -0
  19. metaflow_extensions/outerbounds/plugins/apps/core/config/config_utils.py +868 -0
  20. metaflow_extensions/outerbounds/plugins/apps/core/config/schema_export.py +288 -0
  21. metaflow_extensions/outerbounds/plugins/apps/core/config/typed_configs.py +139 -0
  22. metaflow_extensions/outerbounds/plugins/apps/core/config/typed_init_generator.py +398 -0
  23. metaflow_extensions/outerbounds/plugins/apps/core/config/unified_config.py +1088 -0
  24. metaflow_extensions/outerbounds/plugins/apps/core/config_schema.yaml +337 -0
  25. metaflow_extensions/outerbounds/plugins/apps/core/dependencies.py +115 -0
  26. metaflow_extensions/outerbounds/plugins/apps/core/deployer.py +303 -0
  27. metaflow_extensions/outerbounds/plugins/apps/core/experimental/__init__.py +89 -0
  28. metaflow_extensions/outerbounds/plugins/apps/core/perimeters.py +87 -0
  29. metaflow_extensions/outerbounds/plugins/apps/core/secrets.py +164 -0
  30. metaflow_extensions/outerbounds/plugins/apps/core/utils.py +233 -0
  31. metaflow_extensions/outerbounds/plugins/apps/core/validations.py +17 -0
  32. metaflow_extensions/outerbounds/plugins/aws/assume_role_decorator.py +25 -12
  33. metaflow_extensions/outerbounds/plugins/checkpoint_datastores/coreweave.py +9 -77
  34. metaflow_extensions/outerbounds/plugins/checkpoint_datastores/external_chckpt.py +85 -0
  35. metaflow_extensions/outerbounds/plugins/checkpoint_datastores/nebius.py +7 -78
  36. metaflow_extensions/outerbounds/plugins/fast_bakery/docker_environment.py +6 -2
  37. metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery.py +1 -0
  38. metaflow_extensions/outerbounds/plugins/nvct/nvct_decorator.py +8 -8
  39. metaflow_extensions/outerbounds/plugins/optuna/__init__.py +48 -0
  40. metaflow_extensions/outerbounds/plugins/profilers/simple_card_decorator.py +96 -0
  41. metaflow_extensions/outerbounds/plugins/s3_proxy/__init__.py +7 -0
  42. metaflow_extensions/outerbounds/plugins/s3_proxy/binary_caller.py +132 -0
  43. metaflow_extensions/outerbounds/plugins/s3_proxy/constants.py +11 -0
  44. metaflow_extensions/outerbounds/plugins/s3_proxy/exceptions.py +13 -0
  45. metaflow_extensions/outerbounds/plugins/s3_proxy/proxy_bootstrap.py +59 -0
  46. metaflow_extensions/outerbounds/plugins/s3_proxy/s3_proxy_api.py +93 -0
  47. metaflow_extensions/outerbounds/plugins/s3_proxy/s3_proxy_decorator.py +250 -0
  48. metaflow_extensions/outerbounds/plugins/s3_proxy/s3_proxy_manager.py +225 -0
  49. metaflow_extensions/outerbounds/plugins/snowpark/snowpark_client.py +6 -3
  50. metaflow_extensions/outerbounds/plugins/snowpark/snowpark_decorator.py +13 -7
  51. metaflow_extensions/outerbounds/plugins/snowpark/snowpark_job.py +8 -2
  52. metaflow_extensions/outerbounds/plugins/torchtune/__init__.py +4 -0
  53. metaflow_extensions/outerbounds/plugins/vllm/__init__.py +173 -95
  54. metaflow_extensions/outerbounds/plugins/vllm/status_card.py +9 -9
  55. metaflow_extensions/outerbounds/plugins/vllm/vllm_manager.py +159 -9
  56. metaflow_extensions/outerbounds/remote_config.py +8 -3
  57. metaflow_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.py +62 -1
  58. metaflow_extensions/outerbounds/toplevel/ob_internal.py +2 -0
  59. metaflow_extensions/outerbounds/toplevel/plugins/optuna/__init__.py +1 -0
  60. metaflow_extensions/outerbounds/toplevel/s3_proxy.py +88 -0
  61. {ob_metaflow_extensions-1.1.171rc1.dist-info → ob_metaflow_extensions-1.4.35.dist-info}/METADATA +2 -2
  62. {ob_metaflow_extensions-1.1.171rc1.dist-info → ob_metaflow_extensions-1.4.35.dist-info}/RECORD +64 -22
  63. {ob_metaflow_extensions-1.1.171rc1.dist-info → ob_metaflow_extensions-1.4.35.dist-info}/WHEEL +0 -0
  64. {ob_metaflow_extensions-1.1.171rc1.dist-info → ob_metaflow_extensions-1.4.35.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,288 @@
1
+ """
2
+ Schema Export Module for Unified Configuration System
3
+
4
+ This module provides standalone functions to export configuration schemas in various formats:
5
+ - OpenAPI schemas in YAML or JSON format
6
+ - JSON schemas in YAML or JSON format
7
+
8
+ Usage:
9
+ from schema_export import export_schema, to_openapi_schema, to_json_schema, to_dict
10
+
11
+ # Export schema to file
12
+ export_schema(CoreConfig, "schema.yaml")
13
+ export_schema(CoreConfig, "schema.json", schema_type="json", format="json")
14
+
15
+ # Generate schema in memory
16
+ openapi_schema = to_openapi_schema(CoreConfig)
17
+ json_schema = to_json_schema(CoreConfig)
18
+
19
+ # Export config instance to dict
20
+ config_instance = CoreConfig(name="myapp", port=8000)
21
+ config_dict = to_dict(config_instance)
22
+
23
+ No external dependencies required for basic functionality.
24
+ PyYAML is optional for YAML export support.
25
+ """
26
+
27
+ import json
28
+ import textwrap
29
+ from collections import OrderedDict
30
+ from typing import Any, Dict
31
+
32
+ try:
33
+ import yaml
34
+
35
+ HAS_YAML = True
36
+ except ImportError:
37
+ HAS_YAML = False
38
+
39
+
40
+ def to_openapi_schema(config_class) -> Dict[str, Any]:
41
+ """Generate OpenAPI schema for a configuration class.
42
+
43
+ Args:
44
+ config_class: The configuration class to generate schema for
45
+
46
+ Returns:
47
+ OpenAPI schema dictionary
48
+ """
49
+ return _generate_openapi_schema(config_class)
50
+
51
+
52
+ def to_json_schema(config_class) -> Dict[str, Any]:
53
+ """Generate JSON schema for a configuration class.
54
+
55
+ Args:
56
+ config_class: The configuration class to generate schema for
57
+
58
+ Returns:
59
+ JSON schema dictionary
60
+ """
61
+ return _generate_json_schema(config_class)
62
+
63
+
64
+ def export_schema(
65
+ config_class, filepath: str, schema_type: str = "openapi", format: str = "yaml"
66
+ ) -> None:
67
+ """Export configuration schema to file.
68
+
69
+ Args:
70
+ config_class: The configuration class to export schema for
71
+ filepath: Path to save the schema file
72
+ schema_type: Type of schema to generate ('openapi' or 'json')
73
+ format: Output format ('yaml' or 'json')
74
+
75
+ Examples:
76
+ # Export OpenAPI schema as YAML (default)
77
+ export_schema(CoreConfig, "schema.yaml")
78
+
79
+ # Export JSON schema as YAML
80
+ export_schema(CoreConfig, "schema.yaml", schema_type="json")
81
+
82
+ # Export OpenAPI schema as JSON
83
+ export_schema(CoreConfig, "schema.json", schema_type="openapi", format="json")
84
+
85
+ # Export JSON schema as JSON
86
+ export_schema(CoreConfig, "schema.json", schema_type="json", format="json")
87
+ """
88
+ # Validate inputs
89
+ if schema_type.lower() not in ["openapi", "json"]:
90
+ raise ValueError(
91
+ f"Unsupported schema type: {schema_type}. Use 'openapi' or 'json'."
92
+ )
93
+
94
+ if format.lower() not in ["yaml", "json"]:
95
+ raise ValueError(f"Unsupported format: {format}. Use 'yaml' or 'json'.")
96
+
97
+ # Generate the appropriate schema
98
+ if schema_type.lower() == "openapi":
99
+ base_schema = _generate_openapi_schema(config_class)
100
+ # Wrap in OpenAPI document structure with proper ordering
101
+ schema_data = OrderedDict(
102
+ [
103
+ ("openapi", "3.0.0"),
104
+ (
105
+ "info",
106
+ OrderedDict(
107
+ [
108
+ ("title", f"{config_class.__name__} Configuration Schema"),
109
+ ("version", "1.0.0"),
110
+ ]
111
+ ),
112
+ ),
113
+ (
114
+ "components",
115
+ OrderedDict(
116
+ [
117
+ (
118
+ "schemas",
119
+ OrderedDict([(config_class.__name__, base_schema)]),
120
+ )
121
+ ]
122
+ ),
123
+ ),
124
+ ]
125
+ )
126
+ else: # json schema
127
+ schema_data = _generate_json_schema(config_class)
128
+
129
+ # Export in the requested format
130
+ if format.lower() == "yaml":
131
+ if not HAS_YAML:
132
+ raise ImportError(
133
+ "PyYAML is required for YAML export. Install with: pip install pyyaml"
134
+ )
135
+
136
+ # Custom YAML representer for multiline strings
137
+ def multiline_representer(dumper, data):
138
+ if "\n" in data or len(data) > 80:
139
+ style = "|" # use literal block
140
+ else:
141
+ style = None # normal flow
142
+ return dumper.represent_scalar("tag:yaml.org,2002:str", data, style=style)
143
+
144
+ # Custom YAML representer for OrderedDict to preserve order
145
+ def ordered_dict_representer(dumper, data):
146
+ return dumper.represent_dict(data.items())
147
+
148
+ yaml.add_representer(str, multiline_representer)
149
+ yaml.add_representer(OrderedDict, ordered_dict_representer)
150
+
151
+ with open(filepath, "w") as f:
152
+ yaml.dump(
153
+ schema_data,
154
+ f,
155
+ default_flow_style=False,
156
+ sort_keys=False,
157
+ allow_unicode=True,
158
+ width=120,
159
+ indent=2,
160
+ )
161
+ else: # json format
162
+ with open(filepath, "w") as f:
163
+ json.dump(schema_data, f, indent=2)
164
+
165
+
166
+ # Private helper functions
167
+ def _generate_openapi_schema(cls) -> Dict[str, Any]:
168
+ """Generate OpenAPI schema for a configuration class."""
169
+ # Clean up class docstring for better YAML formatting
170
+ description = f"{cls.__name__} configuration"
171
+ get_description = getattr(cls, "SCHEMA_DOC", None)
172
+ if get_description:
173
+ description = get_description
174
+ elif cls.__doc__:
175
+ # Remove common indentation and clean up whitespace
176
+ cleaned_doc = textwrap.dedent(cls.__doc__).strip()
177
+ # Replace multiple spaces with single spaces but preserve line breaks
178
+ lines = [line.strip() for line in cleaned_doc.split("\n") if line.strip()]
179
+ description = "\n".join(lines)
180
+
181
+ # Create ordered schema with specific order: title, description, type, required, then properties
182
+ schema = OrderedDict(
183
+ [
184
+ ("title", cls.__name__),
185
+ ("description", description),
186
+ ("type", "object"),
187
+ ("required", []),
188
+ ("properties", OrderedDict()),
189
+ ]
190
+ )
191
+
192
+ for field_name, field_info in cls._fields.items():
193
+ if field_name.startswith("_"):
194
+ continue
195
+
196
+ field_schema = _get_field_schema(field_info)
197
+ schema["properties"][field_name] = field_schema
198
+
199
+ # Add to required if field is required
200
+ if field_info.required:
201
+ schema["required"].append(field_name)
202
+
203
+ return schema
204
+
205
+
206
+ def _generate_json_schema(cls) -> Dict[str, Any]:
207
+ """Generate JSON schema for a configuration class."""
208
+ openapi_schema = _generate_openapi_schema(cls)
209
+
210
+ # Create ordered JSON schema with $schema first
211
+ schema = OrderedDict(
212
+ [
213
+ ("$schema", "https://json-schema.org/draft/2020-12/schema"),
214
+ ("title", openapi_schema["title"]),
215
+ ("description", openapi_schema["description"]),
216
+ ("type", openapi_schema["type"]),
217
+ ("required", openapi_schema["required"]),
218
+ ("properties", openapi_schema["properties"]),
219
+ ]
220
+ )
221
+
222
+ return schema
223
+
224
+
225
+ def _get_field_schema(field_info) -> Dict[str, Any]:
226
+ """Generate schema for a single field."""
227
+ field_schema = OrderedDict()
228
+
229
+ # Get description from ConfigField.help first, then CLI metadata
230
+ description = field_info.help or (
231
+ field_info.cli_meta.help if field_info.cli_meta else None
232
+ )
233
+ if description:
234
+ field_schema["description"] = description
235
+
236
+ # Handle field type
237
+ if field_info.field_type == str:
238
+ field_schema["type"] = "string"
239
+ elif field_info.field_type == int:
240
+ field_schema["type"] = "integer"
241
+ elif field_info.field_type == float:
242
+ field_schema["type"] = "number"
243
+ elif field_info.field_type == bool:
244
+ field_schema["type"] = "boolean"
245
+ elif field_info.field_type == list:
246
+ field_schema["type"] = "array"
247
+ field_schema["items"] = {"type": "string"} # Default to string items
248
+ elif field_info.field_type == dict:
249
+ field_schema["type"] = "object"
250
+ field_schema["additionalProperties"] = True
251
+ elif hasattr(field_info.field_type, "_fields"):
252
+ # Nested configuration object
253
+ field_schema = _generate_openapi_schema(field_info.field_type)
254
+ else:
255
+ # Fallback to string for unknown types
256
+ field_schema["type"] = "string"
257
+
258
+ # Add default value
259
+ if field_info.default is not None and not callable(field_info.default):
260
+ field_schema["default"] = field_info.default
261
+
262
+ # Handle choices from CLI metadata
263
+ if field_info.cli_meta and field_info.cli_meta.choices:
264
+ if "type" in field_schema:
265
+ field_schema["enum"] = field_info.cli_meta.choices
266
+
267
+ # Add examples from ConfigField
268
+ if field_info.example is not None:
269
+ field_schema["example"] = field_info.example
270
+
271
+ # Add experimental flag
272
+ if field_info.is_experimental:
273
+ field_schema["experimental"] = True
274
+
275
+ # Add Field Behavior
276
+ if field_info.behavior:
277
+ field_schema["mutation_behavior"] = field_info.behavior
278
+
279
+ if field_info.cli_meta is not None:
280
+ field_schema["cli_option"] = field_info.cli_meta.cli_option_str
281
+
282
+ # Handle validation from CLI metadata
283
+ if field_info.cli_meta and field_info.validation_fn:
284
+ # Add validation hints in description
285
+ if "description" in field_schema:
286
+ field_schema["description"] += " (validation applied)"
287
+
288
+ return field_schema
@@ -0,0 +1,139 @@
1
+ """
2
+ Auto-generated typed classes for ConfigMeta classes.
3
+
4
+ This module provides IDE-friendly typed interfaces for all configuration classes.
5
+ The reason we auto-generate this file is because we want to provide a bridge between what is the ConfigMeta classes and the typed programmatic interface.
6
+ The CoreConfig class is setup in a way that if any additionally params are missed out from being auto-generated then it will not affect the core functionality of the programmatic API.
7
+ The new parameters will just not show up in IDE autocompletions.
8
+ It is fine if this file is not regularly updated by running the script in the .pre-commit-config.app-changes.yaml
9
+ but it is recommended that this file not be deleted or manually edited.
10
+
11
+ """
12
+
13
+ from typing import Optional, List, Dict, Any
14
+ from .unified_config import CoreConfig
15
+
16
+ import sys
17
+ from typing import TYPE_CHECKING
18
+
19
+ # on 3.8+ use the stdlib TypedDict;
20
+ # in TYPE_CHECKING blocks mypy/pyright still pick it up on older Pythons
21
+ if sys.version_info >= (3, 8):
22
+ from typing import TypedDict
23
+ else:
24
+ if TYPE_CHECKING:
25
+ # for the benefit of type-checkers
26
+ from typing import TypedDict # noqa: F401
27
+ # runtime no-op TypedDict shim
28
+ class _TypedDictMeta(type):
29
+ def __new__(cls, name, bases, namespace, total=True):
30
+ # ignore total at runtime
31
+ return super().__new__(cls, name, bases, namespace)
32
+
33
+ class TypedDict(dict, metaclass=_TypedDictMeta):
34
+ # Runtime stand-in for typing.TypedDict on <3.8.
35
+ pass
36
+
37
+
38
+ class ResourceConfigDict(TypedDict, total=False):
39
+ cpu: Optional[str]
40
+ memory: Optional[str]
41
+ gpu: Optional[str]
42
+ disk: Optional[str]
43
+ shared_memory: Optional[str]
44
+
45
+
46
+ class AuthConfigDict(TypedDict, total=False):
47
+ type: Optional[str]
48
+ public: Optional[bool]
49
+
50
+
51
+ class ReplicaConfigDict(TypedDict, total=False):
52
+ fixed: Optional[int]
53
+ min: Optional[int]
54
+ max: Optional[int]
55
+ scaling_policy: Optional["ScalingPolicyConfigDict"]
56
+
57
+
58
+ class ScalingPolicyConfigDict(TypedDict, total=False):
59
+ rpm: Optional[int]
60
+
61
+
62
+ class DependencyConfigDict(TypedDict, total=False):
63
+ from_requirements_file: Optional[str]
64
+ from_pyproject_toml: Optional[str]
65
+ python: Optional[str]
66
+ pypi: Optional[dict]
67
+ conda: Optional[dict]
68
+
69
+
70
+ class PackageConfigDict(TypedDict, total=False):
71
+ src_paths: Optional[list]
72
+ suffixes: Optional[list]
73
+
74
+
75
+ class TypedCoreConfig:
76
+ def __init__(
77
+ self,
78
+ name: Optional[str] = None,
79
+ port: Optional[int] = None,
80
+ description: Optional[str] = None,
81
+ app_type: Optional[str] = None,
82
+ image: Optional[str] = None,
83
+ tags: Optional[list] = None,
84
+ secrets: Optional[list] = None,
85
+ compute_pools: Optional[list] = None,
86
+ environment: Optional[dict] = None,
87
+ commands: Optional[list] = None,
88
+ resources: Optional[ResourceConfigDict] = None,
89
+ auth: Optional[AuthConfigDict] = None,
90
+ replicas: Optional[ReplicaConfigDict] = None,
91
+ dependencies: Optional[DependencyConfigDict] = None,
92
+ package: Optional[PackageConfigDict] = None,
93
+ no_deps: Optional[bool] = None,
94
+ force_upgrade: Optional[bool] = None,
95
+ persistence: Optional[str] = None,
96
+ project: Optional[str] = None,
97
+ branch: Optional[str] = None,
98
+ models: Optional[list] = None,
99
+ data: Optional[list] = None,
100
+ generate_static_url: Optional[bool] = None,
101
+ **kwargs
102
+ ) -> None:
103
+ self._kwargs = {
104
+ "name": name,
105
+ "port": port,
106
+ "description": description,
107
+ "app_type": app_type,
108
+ "image": image,
109
+ "tags": tags,
110
+ "secrets": secrets,
111
+ "compute_pools": compute_pools,
112
+ "environment": environment,
113
+ "commands": commands,
114
+ "resources": resources,
115
+ "auth": auth,
116
+ "replicas": replicas,
117
+ "dependencies": dependencies,
118
+ "package": package,
119
+ "no_deps": no_deps,
120
+ "force_upgrade": force_upgrade,
121
+ "persistence": persistence,
122
+ "project": project,
123
+ "branch": branch,
124
+ "models": models,
125
+ "data": data,
126
+ "generate_static_url": generate_static_url,
127
+ }
128
+ # Add any additional kwargs
129
+ self._kwargs.update(kwargs)
130
+ # Remove None values
131
+ self._kwargs = {k: v for k, v in self._kwargs.items() if v is not None}
132
+ self._config_class = CoreConfig
133
+ self._config = self.create_config()
134
+
135
+ def create_config(self) -> CoreConfig:
136
+ return CoreConfig.from_dict(self._kwargs)
137
+
138
+ def to_dict(self) -> Dict[str, Any]:
139
+ return self._config.to_dict()