[Mlir-commits] [mlir] [MLIR][Python] remove PyYAML as a dep (PR #169145)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-linalg

Author: Maksim Levental (makslevental)

<details>
<summary>Changes</summary>

PyYAML is not an actual use-time/runtime dependency of our bindings. It is only used during build time - during `mlir-linalg-ods-gen`: https://github.com/llvm/llvm-project/blob/93097b2d47c87bf5eee0a2612d961c7a01831eab/mlir/tools/mlir-linalg-ods-gen/update_core_linalg_named_ops.sh.in#L29

This PR does the minimal refactor to remove the need during actual run/use time. 

~~Since we do actually have a test for `linalg.opdsl.dump_oplib` in-tree, the way I'm testing this PR is by temporarily removing PyYAML from `.ci/all_requirements.txt` (i.e., not installing during pre-commit) and showing that all other tests pass.~~ 

It turns out, even without explicitly installing, PyYAML is somehow installed (probably a transitive dep of one of the other deps in `.ci/all_requirements.txt`). Anyway, I've tested locally and downstream and anyone else can repro by doing `pip uninstall pyyaml` and then `ninja check-mlir-python` with this branch.

Note, I would prefer to also split `mlir/python/requirements.txt` into a pair of requirements files (one which contains only the runtime deps and one which contains both the runtime and build deps) but probably that would break all of the builders so I'm leaving it as is.

---
Full diff: https://github.com/llvm/llvm-project/pull/169145.diff


1 Files Affected:

- (modified) mlir/python/mlir/dialects/linalg/opdsl/lang/yaml_helper.py (+43-23) 


``````````diff
diff --git a/mlir/python/mlir/dialects/linalg/opdsl/lang/yaml_helper.py b/mlir/python/mlir/dialects/linalg/opdsl/lang/yaml_helper.py
index 1672656b3a1f8..f3ad92d782eff 100644
--- a/mlir/python/mlir/dialects/linalg/opdsl/lang/yaml_helper.py
+++ b/mlir/python/mlir/dialects/linalg/opdsl/lang/yaml_helper.py
@@ -4,23 +4,30 @@
 """YAML serialization is routed through here to centralize common logic."""
 
 import sys
+import warnings
+
+
+def multiline_str_representer(dumper, data):
+    if len(data.splitlines()) > 1:
+        return dumper.represent_scalar("tag:yaml.org,2002:str", data, style="|")
+    else:
+        return dumper.represent_scalar("tag:yaml.org,2002:str", data)
+
 
 try:
-    import yaml
+    from yaml import YAMLObject, add_representer
+
+    add_representer(str, multiline_str_representer)
 except ModuleNotFoundError as e:
-    raise ModuleNotFoundError(
-        f"This tool requires PyYAML but it was not installed. "
-        f"Recommend: {sys.executable} -m pip install PyYAML"
-    ) from e
+    warnings.warn(
+        "PyYAML is not installed; you will not be able to generate linalg yaml definitions."
+    )
 
-__all__ = [
-    "yaml_dump",
-    "yaml_dump_all",
-    "YAMLObject",
-]
+    class YAMLObject:
+        pass
 
 
-class YAMLObject(yaml.YAMLObject):
+class YAMLObject(YAMLObject):
     @classmethod
     def to_yaml(cls, dumper, self):
         """Default to a custom dictionary mapping."""
@@ -33,21 +40,34 @@ def as_linalg_yaml(self):
         return yaml_dump(self)
 
 
-def multiline_str_representer(dumper, data):
-    if len(data.splitlines()) > 1:
-        return dumper.represent_scalar("tag:yaml.org,2002:str", data, style="|")
-    else:
-        return dumper.represent_scalar("tag:yaml.org,2002:str", data)
+def yaml_dump(data, sort_keys=False, **kwargs):
+    try:
+        import yaml
 
+        return yaml.dump(data, sort_keys=sort_keys, **kwargs)
+    except ModuleNotFoundError as e:
+        raise ModuleNotFoundError(
+            f"This tool requires PyYAML but it was not installed. "
+            f"Recommend: {sys.executable} -m pip install PyYAML"
+        ) from e
 
-yaml.add_representer(str, multiline_str_representer)
 
+def yaml_dump_all(data, sort_keys=False, explicit_start=True, **kwargs):
+    try:
+        import yaml
 
-def yaml_dump(data, sort_keys=False, **kwargs):
-    return yaml.dump(data, sort_keys=sort_keys, **kwargs)
+        return yaml.dump_all(
+            data, sort_keys=sort_keys, explicit_start=explicit_start, **kwargs
+        )
+    except ModuleNotFoundError as e:
+        raise ModuleNotFoundError(
+            f"This tool requires PyYAML but it was not installed. "
+            f"Recommend: {sys.executable} -m pip install PyYAML"
+        ) from e
 
 
-def yaml_dump_all(data, sort_keys=False, explicit_start=True, **kwargs):
-    return yaml.dump_all(
-        data, sort_keys=sort_keys, explicit_start=explicit_start, **kwargs
-    )
+__all__ = [
+    "yaml_dump",
+    "yaml_dump_all",
+    "YAMLObject",
+]

``````````

</details>


https://github.com/llvm/llvm-project/pull/169145