flyteplugins-omegaconf 2.1.6__py3-none-any.whl

flyteplugins/omegaconf/__init__.py

@@ -0,0 +1,24 @@
+"""OmegaConf DictConfig/ListConfig support for Flyte."""
+
+from __future__ import annotations
+
+import functools
+
+from flyte.types._type_engine import TypeEngine
+
+from .base_transformer import DictConfigTransformer, ListConfigTransformer
+
+
+@functools.lru_cache(maxsize=None)
+def register_omegaconf_transformers() -> None:
+    """Register OmegaConf transformers with Flyte TypeEngine.
+
+    Called via the ``flyte.plugins.types`` entry point on import, or manually
+    by importing this package.
+    """
+    TypeEngine.register(DictConfigTransformer())
+    TypeEngine.register(ListConfigTransformer())
+
+
+# Register at module import time for backwards compatibility.
+register_omegaconf_transformers()

flyteplugins/omegaconf/base_transformer.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+from typing import Generic, Type, TypeVar
+
+import msgpack
+from flyte.types._type_engine import (
+    CACHE_KEY_METADATA,
+    MESSAGEPACK,
+    SERIALIZATION_FORMAT,
+    TypeTransformer,
+)
+from flyteidl2.core.literals_pb2 import Binary, Literal, Scalar
+from flyteidl2.core.types_pb2 import LiteralType, SimpleType, TypeAnnotation
+from google.protobuf import struct_pb2
+
+from omegaconf import DictConfig, ListConfig
+
+from .codec import deserialize_omegaconf, serialize_omegaconf
+
+T = TypeVar("T", DictConfig, ListConfig)
+
+
+class OmegaConfTransformerBase(TypeTransformer[T], Generic[T]):
+    def __init__(self, name: str, container_type: Type[T]):
+        super().__init__(name, container_type)
+        self._container_type = container_type
+
+    def get_literal_type(self, t: Type[T]) -> LiteralType:
+        meta_struct = struct_pb2.Struct()
+        meta_struct.update({CACHE_KEY_METADATA: {SERIALIZATION_FORMAT: MESSAGEPACK}})
+        return LiteralType(
+            simple=SimpleType.STRUCT,
+            annotation=TypeAnnotation(annotations=meta_struct),
+        )
+
+    async def to_literal(
+        self,
+        python_val: T,
+        python_type: Type[T],
+        expected: LiteralType,
+    ) -> Literal:
+        payload = serialize_omegaconf(python_val)
+        msgpack_bytes = msgpack.dumps(payload)
+        return Literal(scalar=Scalar(binary=Binary(value=msgpack_bytes, tag=MESSAGEPACK)))
+
+    def from_binary_idl(self, binary_idl_object: Binary, expected_python_type: Type[T]) -> T:
+        if binary_idl_object.tag != MESSAGEPACK:
+            raise TypeError(f"Unsupported binary format: `{binary_idl_object.tag}`")
+        payload = msgpack.loads(binary_idl_object.value, strict_map_key=False)
+        config = deserialize_omegaconf(payload)
+        if not isinstance(config, self._container_type):
+            raise TypeError(f"Expected {self._container_type.__name__} payload, got {type(config).__name__}")
+        return config
+
+    async def to_python_value(self, lv: Literal, expected_python_type: Type[T]) -> T:
+        if lv and lv.HasField("scalar") and lv.scalar.HasField("binary"):
+            return self.from_binary_idl(lv.scalar.binary, expected_python_type)
+        raise TypeError(f"Cannot convert literal to {self._container_type.__name__}: {lv}")
+
+
+class DictConfigTransformer(OmegaConfTransformerBase[DictConfig]):
+    def __init__(self):
+        super().__init__("OmegaConf DictConfig Transformer", DictConfig)
+
+
+class ListConfigTransformer(OmegaConfTransformerBase[ListConfig]):
+    def __init__(self):
+        super().__init__("OmegaConf ListConfig Transformer", ListConfig)

flyteplugins/omegaconf/codec.py

@@ -0,0 +1,169 @@
+from __future__ import annotations
+
+import importlib
+from enum import Enum
+from pathlib import Path, PurePath
+from typing import Any
+
+from omegaconf import MISSING, DictConfig, ListConfig, OmegaConf
+
+PAYLOAD_MARKER = "__flyte_omegaconf__"
+PAYLOAD_KIND = "kind"
+PAYLOAD_VALUES = "values"
+PAYLOAD_SCHEMA = "schema"
+PAYLOAD_TYPE = "type"
+PAYLOAD_NAME = "name"
+PAYLOAD_VALUE = "value"
+
+KIND_DICT = "dict"
+KIND_LIST = "list"
+KIND_MISSING = "missing"
+KIND_ENUM = "enum"
+KIND_PATH = "path"
+KIND_TUPLE = "tuple"
+
+
+def serialize_omegaconf(node: DictConfig | ListConfig) -> dict[str, Any]:
+    return _serialize_value(node)
+
+
+def deserialize_omegaconf(payload: dict[str, Any]) -> DictConfig | ListConfig:
+    config = _deserialize_value(payload)
+    if isinstance(config, (DictConfig, ListConfig)):
+        return config
+    raise TypeError(f"Expected OmegaConf container payload, got {type(config).__name__}")
+
+
+def _serialize_value(value: Any) -> Any:
+    if isinstance(value, DictConfig):
+        payload: dict[str, Any] = {
+            PAYLOAD_MARKER: True,
+            PAYLOAD_KIND: KIND_DICT,
+            PAYLOAD_VALUES: {},
+        }
+        schema_name = _schema_name(value)
+        if schema_name is not None:
+            payload[PAYLOAD_SCHEMA] = schema_name
+
+        for key in value.keys():
+            child_node = value._get_node(key)
+            payload[PAYLOAD_VALUES][key] = _serialize_child(value, key, child_node)
+        return payload
+
+    if isinstance(value, ListConfig):
+        payload = {
+            PAYLOAD_MARKER: True,
+            PAYLOAD_KIND: KIND_LIST,
+            PAYLOAD_VALUES: [],
+        }
+        for index in range(len(value)):
+            child_node = value._get_node(index)
+            payload[PAYLOAD_VALUES].append(_serialize_child(value, index, child_node))
+        return payload
+
+    if isinstance(value, Enum):
+        return {
+            PAYLOAD_MARKER: True,
+            PAYLOAD_KIND: KIND_ENUM,
+            PAYLOAD_TYPE: _qualified_name(type(value)),
+            PAYLOAD_NAME: value.name,
+            PAYLOAD_VALUE: value.value,
+        }
+
+    if isinstance(value, PurePath):
+        return {
+            PAYLOAD_MARKER: True,
+            PAYLOAD_KIND: KIND_PATH,
+            PAYLOAD_VALUE: str(value),
+        }
+
+    if isinstance(value, tuple):
+        return {
+            PAYLOAD_MARKER: True,
+            PAYLOAD_KIND: KIND_TUPLE,
+            PAYLOAD_VALUES: [_serialize_value(item) for item in value],
+        }
+
+    if isinstance(value, dict):
+        return {key: _serialize_value(item) for key, item in value.items()}
+
+    if isinstance(value, list):
+        return [_serialize_value(item) for item in value]
+
+    return value
+
+
+def _serialize_child(parent: DictConfig | ListConfig, key: str | int, child_node: Any) -> Any:
+    if child_node._is_missing():
+        return {
+            PAYLOAD_MARKER: True,
+            PAYLOAD_KIND: KIND_MISSING,
+        }
+    return _serialize_value(parent[key])
+
+
+def _deserialize_value(payload: Any) -> Any:
+    if isinstance(payload, list):
+        return [_deserialize_value(item) for item in payload]
+
+    if not isinstance(payload, dict):
+        return payload
+
+    if payload.get(PAYLOAD_MARKER) is not True:
+        return {key: _deserialize_value(value) for key, value in payload.items()}
+
+    kind = payload[PAYLOAD_KIND]
+
+    if kind == KIND_MISSING:
+        return MISSING
+
+    if kind == KIND_ENUM:
+        enum_value = payload[PAYLOAD_VALUE]
+        enum_name = payload[PAYLOAD_NAME]
+        enum_type_name = payload[PAYLOAD_TYPE]
+        try:
+            enum_type = _import_class(enum_type_name)
+            return enum_type[enum_name]
+        except (ImportError, AttributeError, ModuleNotFoundError, KeyError, ValueError):
+            return enum_value
+
+    if kind == KIND_PATH:
+        return Path(payload[PAYLOAD_VALUE])
+
+    if kind == KIND_TUPLE:
+        return tuple(_deserialize_value(item) for item in payload[PAYLOAD_VALUES])
+
+    if kind == KIND_LIST:
+        values = [_deserialize_value(item) for item in payload[PAYLOAD_VALUES]]
+        return OmegaConf.create(values)
+
+    if kind == KIND_DICT:
+        values = {key: _deserialize_value(value) for key, value in payload[PAYLOAD_VALUES].items()}
+        schema_name = payload.get(PAYLOAD_SCHEMA)
+        if schema_name is None:
+            return OmegaConf.create(values)
+
+        try:
+            schema_type = _import_class(schema_name)
+            return OmegaConf.merge(OmegaConf.structured(schema_type), values)
+        except (ImportError, AttributeError, ModuleNotFoundError):
+            return OmegaConf.create(values)
+
+    raise TypeError(f"Unsupported OmegaConf payload kind: {kind}")
+
+
+def _schema_name(config: DictConfig) -> str | None:
+    schema_type = OmegaConf.get_type(config)
+    if schema_type in (None, dict):
+        return None
+    return _qualified_name(schema_type)
+
+
+def _qualified_name(value_type: type) -> str:
+    return f"{value_type.__module__}.{value_type.__qualname__}"
+
+
+def _import_class(fully_qualified_name: str) -> type:
+    module_name, class_name = fully_qualified_name.rsplit(".", 1)
+    module = importlib.import_module(module_name)
+    return getattr(module, class_name)

flyteplugins_omegaconf-2.1.6.dist-info/METADATA

@@ -0,0 +1,193 @@
+Metadata-Version: 2.4
+Name: flyteplugins-omegaconf
+Version: 2.1.6
+Summary: OmegaConf DictConfig/ListConfig support for Flyte
+Author-email: Samhita Alla <samhita@union.ai>
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: flyte
+Requires-Dist: omegaconf
+
+# flyteplugins-omegaconf
+
+Enables [OmegaConf](https://omegaconf.readthedocs.io/) `DictConfig` and `ListConfig` as typed inputs and outputs for Flyte tasks.
+
+## Installation
+
+```bash
+pip install flyteplugins-omegaconf
+```
+
+Installing the package automatically registers the `DictConfig` and `ListConfig` transformers with Flyte's TypeEngine via the `flyte.plugins.types` entry point.
+
+## Usage
+
+### DictConfig as task inputs and outputs
+
+```python
+import flyte
+from omegaconf import DictConfig, OmegaConf
+
+env = flyte.TaskEnvironment(name="training", image=...)
+
+@env.task
+async def preprocess(cfg: DictConfig) -> DictConfig:
+    return OmegaConf.merge(cfg, {"data": {"normalized": True}})
+
+@env.task
+async def train(cfg: DictConfig) -> float:
+    return run_experiment(cfg.optimizer.lr, cfg.training.epochs)
+
+@env.task
+async def pipeline() -> float:
+    cfg = OmegaConf.create({"optimizer": {"lr": 0.001}, "training": {"epochs": 10}})
+    processed = await preprocess(cfg)
+    return await train(processed)
+```
+
+### ListConfig as task inputs and outputs
+
+```python
+from omegaconf import ListConfig, OmegaConf
+
+@env.task
+async def build_lr_schedule(base_lr: float, num_stages: int) -> ListConfig:
+    return OmegaConf.create([base_lr * (0.5 ** i) for i in range(num_stages)])
+
+@env.task
+async def train_with_schedule(cfg: DictConfig, lr_schedule: ListConfig) -> float:
+    final_lr = float(lr_schedule[-1])
+    ...
+```
+
+## Ways to construct a DictConfig
+
+All of the following are valid ways to create a `DictConfig` to pass to a task:
+
+### 1. From a plain dict
+
+```python
+cfg = OmegaConf.create({"optimizer": {"lr": 0.001}, "training": {"epochs": 10}})
+flyte.run(train, cfg=cfg)
+```
+
+### 2. From a YAML file
+
+```python
+cfg = OmegaConf.load("config.yaml")
+flyte.run(train, cfg=cfg)
+```
+
+### 3. From a typed dataclass (structured config)
+
+```python
+from dataclasses import dataclass, field
+from omegaconf import OmegaConf
+
+@dataclass
+class OptimizerConf:
+    lr: float = 0.001
+    weight_decay: float = 1e-4
+
+@dataclass
+class TrainConf:
+    optimizer: OptimizerConf = field(default_factory=OptimizerConf)
+    epochs: int = 10
+
+cfg = OmegaConf.structured(TrainConf())
+flyte.run(train, cfg=cfg)
+```
+
+Structured configs provide **type validation at assignment time**: `cfg.optimizer.lr = "oops"` raises `omegaconf.ValidationError`.
+
+### 4. Merging base config with overrides
+
+```python
+base = OmegaConf.load("config.yaml")
+override = OmegaConf.create({"optimizer": {"lr": 0.01}})
+cfg = OmegaConf.merge(base, override)
+flyte.run(train, cfg=cfg)
+```
+
+### 5. Structured config with MISSING required fields
+
+```python
+from omegaconf import MISSING
+
+@dataclass
+class TrainConf:
+    data_path: str = MISSING  # must be set before accessing
+    epochs: int = 10
+
+# Pass with MISSING still unset — serialization succeeds
+cfg = OmegaConf.structured(TrainConf())
+flyte.run(train, cfg=cfg)
+
+# Or fill it before passing
+cfg = OmegaConf.structured(TrainConf(data_path="/data/imagenet"))
+flyte.run(train, cfg=cfg)
+```
+
+A config with an unset `MISSING` field serializes and deserializes successfully — the `MISSING` sentinel is preserved through the wire format. Accessing the field raises `MissingMandatoryValue`, so the task will fail if it tries to read an unfilled field.
+
+## Structured config deserialization
+
+When a `DictConfig` is deserialized in a receiving task, the plugin uses **Auto mode**: it attempts to reconstruct the original dataclass-backed config, and falls back to a plain `DictConfig` if the class is not importable in the receiving task's environment.
+
+```python
+# Task A produces a structured config
+cfg = OmegaConf.structured(TrainConf(lr=0.01))
+# serialized payload: {"base_dataclass": "mymodule.TrainConf", "values": {...}}
+
+# Task B receives it
+async def task_b(cfg: DictConfig) -> ...:
+    # If TrainConf is importable: cfg is a TrainConf-backed DictConfig (type-validated)
+    # If TrainConf is not importable: cfg is a plain DictConfig (no schema)
+    OmegaConf.get_type(cfg)  # TrainConf or dict
+```
+
+To ensure structured configs survive task hops, make sure the dataclass is defined in a module importable by all tasks in the pipeline.
+
+## Wire format
+
+Both `DictConfig` and `ListConfig` are serialized as MessagePack binaries with tag `"msgpack"`:
+
+```
+Literal(scalar=Scalar(binary=Binary(value=<msgpack bytes>, tag="msgpack")))
+```
+
+**DictConfig payload** (msgpack-encoded dict):
+
+```json
+{
+  "base_dataclass": "mymodule.TrainConf",
+  "values": { "optimizer": { "lr": 0.001 }, "training": { "epochs": 10 } }
+}
+```
+
+For plain dict-backed configs, `base_dataclass` is `"builtins.dict"`.
+
+**ListConfig payload** (msgpack-encoded list):
+
+```json
+[0.001, 0.01, 0.1]
+```
+
+OmegaConf variable interpolations are **resolved** at serialization time (`resolve=True`). The wire representation always contains concrete values.
+
+## Limitations
+
+- **Structured config schema strictness**: merging keys that don't exist as dataclass fields raises an error. Only declare structured configs when all possible keys are known upfront.
+- **`MISSING` fields**: a `DictConfig` with unset `MISSING` fields serializes fine — the sentinel is preserved on the wire and accessing it still raises `MissingMandatoryValue`. However, in plain dict mode (when the originating dataclass is not importable in the receiving task), the field's type annotation is lost: the node becomes an `AnyNode` instead of the declared type (e.g. `StringNode`). In Auto mode, the schema is recovered from the dataclass, so the annotation is preserved.
+- **ListConfig structured configs**: `ListConfig` always round-trips as a plain `ListConfig` — there is no structured (typed-element) ListConfig support.
+- **Key types**: OmegaConf enforces string keys for `DictConfig`; integer-keyed dicts are not supported.
+- **Class importability**: structured config reconstruction requires the dataclass to be importable in the receiving task. If it is not, the config falls back to a plain `DictConfig` (Auto mode).
+
+## Examples
+
+See the [`examples/`](examples/) directory:
+
+- [`example_dictconfig.py`](examples/example_dictconfig.py) — plain dict configs, nested, interpolation, merging
+- [`example_structured_config.py`](examples/example_structured_config.py) — structured configs, type validation, MISSING fields, config resolution
+- [`example_listconfig.py`](examples/example_listconfig.py) — numeric lists, nested lists, list of dicts, LR schedules
+- [`example_pipeline.py`](examples/example_pipeline.py) — multi-task pipeline with DictConfig and ListConfig flowing between tasks

flyteplugins_omegaconf-2.1.6.dist-info/RECORD

@@ -0,0 +1,8 @@
+flyteplugins/omegaconf/__init__.py,sha256=NZcB8g3E1FUspUWh8fqdtdWGTen85qgTHm1bHfsev5o,696
+flyteplugins/omegaconf/base_transformer.py,sha256=zPz2hMcH1mFlJCYD_iZS2g34GBv0LE3EtF4ZITCtA3w,2642
+flyteplugins/omegaconf/codec.py,sha256=CtGCZl7s04weRXXQkKUDdE1XC51AO9CpAj3Zgq_5ksI,5183
+flyteplugins_omegaconf-2.1.6.dist-info/METADATA,sha256=R6T6dyJvdGb1CJyGNBtVjYSgY-yxHgHHP3c_QJPiaJo,6838
+flyteplugins_omegaconf-2.1.6.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+flyteplugins_omegaconf-2.1.6.dist-info/entry_points.txt,sha256=F992e_5r188DCSEp6CBXmVL1gMUQ6EXrw9R_7rhAB38,89
+flyteplugins_omegaconf-2.1.6.dist-info/top_level.txt,sha256=cgd779rPu9EsvdtuYgUxNHHgElaQvPn74KhB5XSeMBE,13
+flyteplugins_omegaconf-2.1.6.dist-info/RECORD,,

flyteplugins_omegaconf-2.1.6.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

flyteplugins_omegaconf-2.1.6.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[flyte.plugins.types]
+omegaconf = flyteplugins.omegaconf:register_omegaconf_transformers

flyteplugins_omegaconf-2.1.6.dist-info/top_level.txt

@@ -0,0 +1 @@
+flyteplugins