Use single REGISTRY

The register method and issue manager factories will be more more maintainable if every issue manage is registered to a task. For now, only label issue checks work for regression, as that is the only issue manger registered for regression.
cleanlab · elisno · Nov 20, 2023 · Aug 7, 2023 · Aug 7, 2023 · Aug 7, 2023
commit 585d14a9e16df8bbad5268861192cd90f2f1a587
diff --git a/cleanlab/datalab/internal/issue_manager_factory.py b/cleanlab/datalab/internal/issue_manager_factory.py
@@ -50,16 +50,14 @@
 )
 from cleanlab.datalab.internal.issue_manager.regression import RegressionLabelIssueManager
 
-REGISTRY: Dict[str, Type[IssueManager]] = {
-    "outlier": OutlierIssueManager,
-    "near_duplicate": NearDuplicateIssueManager,
-    "non_iid": NonIIDIssueManager,
-}
 
-TASK_SPECIFIC_REGISTRY: Dict[str, Dict[str, Type[IssueManager]]] = {
+REGISTRY: Dict[str, Dict[str, Type[IssueManager]]] = {
     "classification": {
         "label": LabelIssueManager,
         "class_imbalance": ClassImbalanceIssueManager,
+        "outlier": OutlierIssueManager,
+        "near_duplicate": NearDuplicateIssueManager,
+        "non_iid": NonIIDIssueManager,
     },
     "regression": {
         "label": RegressionLabelIssueManager,
@@ -94,27 +92,22 @@ def from_str(cls, issue_type: str, task: str) -> Type[IssueManager]:
                 "issue_type must be a string, not a list. Try using from_list instead."
             )
 
-        if task is None and issue_type not in REGISTRY:
-            raise ValueError(f"Invalid issue type: {issue_type}")
-        if task not in TASK_SPECIFIC_REGISTRY:
+        if task not in REGISTRY:
             raise ValueError(
-                f"Invalid task type: {task}, must be in {list(TASK_SPECIFIC_REGISTRY.keys())}"
+                f"Invalid task type: {task}, must be in {list(REGISTRY.keys())}"
             )
-        if issue_type not in TASK_SPECIFIC_REGISTRY[task] and issue_type not in REGISTRY:
+        if issue_type not in REGISTRY[task]:
             raise ValueError(f"Invalid issue type: {issue_type} for task {task}")
 
-        if issue_type in TASK_SPECIFIC_REGISTRY[task]:
-            return TASK_SPECIFIC_REGISTRY[task][issue_type]
-
-        return REGISTRY[issue_type]
+        return REGISTRY[task][issue_type]
 
     @classmethod
     def from_list(cls, issue_types: List[str], task: str) -> List[Type[IssueManager]]:
         """Constructs a list of concrete issue manager classes from a list of strings."""
         return [cls.from_str(issue_type, task) for issue_type in issue_types]
 
 
-def register(cls: Type[IssueManager], task: str) -> Type[IssueManager]:
+def register(cls: Type[IssueManager], task: str="classification") -> Type[IssueManager]:
     """Registers the issue manager factory.
 
     Parameters
@@ -171,24 +164,18 @@ def find_issues(self, **kwargs):
 
     name: str = str(cls.issue_name)
 
-    if task is not None and task not in TASK_SPECIFIC_REGISTRY:
+    if task not in REGISTRY:
         raise ValueError(
-            f"Invalid task type: {task}, must be in {list(TASK_SPECIFIC_REGISTRY.keys())}"
+            f"Invalid task type: {task}, must be in {list(REGISTRY.keys())}"
         )
 
-    if task is not None and name in TASK_SPECIFIC_REGISTRY[task]:
+    if name in REGISTRY[task]:
         print(
             f"Warning: Overwriting existing issue manager {name} with {cls} for task {task}."
             "This may cause unexpected behavior."
         )
 
-    if task is None and name in REGISTRY:
-        print(
-            f"Warning: Overwriting existing issue manager {name} with {cls}."
-            "This may cause unexpected behavior."
-        )
-
-    REGISTRY[name] = cls
+    REGISTRY[task][name] = cls
     return cls
 
 
@@ -201,7 +188,7 @@ def list_possible_issue_types(task: str) -> List[str]:
     --------
     :py:class:`REGISTRY <cleanlab.datalab.internal.issue_manager_factory.REGISTRY>` : All available issue types and their corresponding issue managers can be found here.
     """
-    return list(REGISTRY.keys()) + list(TASK_SPECIFIC_REGISTRY.get(task, []))
+    return list(REGISTRY.get(task, []))
 
 
 def list_default_issue_types(task: str) -> List[str]:
@@ -213,7 +200,7 @@ def list_default_issue_types(task: str) -> List[str]:
     :py:class:`REGISTRY <cleanlab.datalab.internal.issue_manager_factory.REGISTRY>` : All available issue types and their corresponding issue managers can be found here.
     """
     if task == "regression":
-        default_issue_types = ["label", "outlier", "near_duplicate", "non_iid"]
+        default_issue_types = ["label"]
     else:
         default_issue_types = [
             "label",

diff --git a/tests/datalab/test_datalab.py b/tests/datalab/test_datalab.py
@@ -624,7 +624,7 @@ def test_custom_issue_manager_registered(self, lab, custom_issue_manager):
         """Test that a custom issue manager that is registered will be used."""
         from cleanlab.datalab.internal.issue_manager_factory import register
 
-        register(custom_issue_manager, task=None)
+        register(custom_issue_manager)
 
         assert lab.issues.empty
         assert lab.issue_summary.empty
@@ -648,7 +648,7 @@ def test_find_issues_for_custom_issue_manager_with_custom_kwarg(
         """Test that a custom issue manager that is registered will be used."""
         from cleanlab.datalab.internal.issue_manager_factory import register
 
-        register(custom_issue_manager, task=None)
+        register(custom_issue_manager)
 
         assert lab.issues.empty
         assert lab.issue_summary.empty
@@ -668,7 +668,8 @@ def test_find_issues_for_custom_issue_manager_with_custom_kwarg(
         # Clean up registry
         from cleanlab.datalab.internal.issue_manager_factory import REGISTRY
 
-        REGISTRY.pop(custom_issue_manager.issue_name)
+        # Find the custom issue manager in the registry and remove it
+        REGISTRY["classification"].pop(custom_issue_manager.issue_name)
 
 
 @pytest.mark.parametrize(
@@ -910,6 +911,10 @@ def test_regression():
     test_df = pd.DataFrame(X, columns=["c1", "c2", "c3"])
     test_df["y"] = y
     lab = Datalab(data=test_df, label_name="y", task="regression")
+
+    assert set(lab.list_default_issue_types()) == set(["label"])
+    assert set(lab.list_possible_issue_types()) == set(["label"])
+
     lab.find_issues(features=X)
     lab.report()
     summary = lab.get_issue_summary()

diff --git a/tests/datalab/test_factory.py b/tests/datalab/test_factory.py
@@ -28,12 +28,12 @@ def test_list_possible_issue_types(registry):
     class TestIssueManager(IssueManager):
         issue_name = test_key
 
-    TestIssueManager = register(TestIssueManager, task=None)
+    TestIssueManager = register(TestIssueManager)
 
     issue_types = lab.list_possible_issue_types()
     assert set(issue_types) == set(
         possible_issues + [test_key]
     ), "New issue type should be added to the list"
 
     # Clean up
-    del registry[test_key]
+    del registry["classification"][test_key]
diff --git a/tests/datalab/test_issue_finder.py b/tests/datalab/test_issue_finder.py
@@ -94,6 +94,7 @@ def test_get_available_issue_types(self, issue_finder):
             {"label": {"some_arg": "some_value"}, "outlier": {}},
             {},
         ]
+        supported_issue_types = ["label"]
         for issue_types in issue_types_dicts:
             available_issue_types = issue_finder.get_available_issue_types(issue_types=issue_types)
             fail_msg = f"Failed to get available issue types with issue_types={issue_types}"
@@ -103,4 +104,4 @@ def test_get_available_issue_types(self, issue_finder):
         kwargs = {k: k for k in ["pred_probs", "features", "knn_graph"]}
         kwargs["issue_types"] = {"label": {}}
         available_issue_types = issue_finder.get_available_issue_types(**kwargs)
-        assert available_issue_types == {"label": {}}
+        assert available_issue_types == {"label": {"features": "features"}}
diff --git a/tests/datalab/test_issue_manager.py b/tests/datalab/test_issue_manager.py
@@ -48,10 +48,9 @@ class Foo(IssueManager):
         def find_issues(self):
             pass
 
-    Foo = register(Foo, task=None)
+    Foo = register(Foo)
 
-    assert "foo" in REGISTRY
-    assert REGISTRY["foo"] == Foo
+    assert REGISTRY["classification"].get("foo") == Foo
 
     # Reregistering should overwrite the existing class, put print a warning
 
@@ -63,14 +62,17 @@ class NewFoo(IssueManager):
         def find_issues(self):
             pass
 
-    NewFoo = register(NewFoo, task=None)
+    NewFoo = register(NewFoo)
 
-    assert "foo" in REGISTRY
-    assert REGISTRY["foo"] == NewFoo
+    assert REGISTRY["classification"].get("foo") == NewFoo
     assert all(
         [
             text in sys.stdout.getvalue()
-            for text in ["Warning: Overwriting existing issue manager foo with ", "NewFoo"]
+            for text in [
+                "Warning: Overwriting existing issue manager foo with ",
+                "NewFoo",
+                " for task classification.",
+            ]
         ]
     ), "Should print a warning"
 
@@ -83,8 +85,7 @@ def find_issues(self):
 
     NewerFoo = register(NewerFoo, task="classification")
 
-    assert "label" in REGISTRY
-    assert REGISTRY["label"] == NewerFoo
+    assert REGISTRY["classification"].get("label") == NewerFoo
     assert all(
         [
             text in sys.stdout.getvalue()
@@ -95,3 +96,14 @@ def find_issues(self):
             ]
         ]
     ), "Should print a warning"
+
+    # Registering any issue manager for another task is permitted
+    class Bar(IssueManager):
+        issue_name = "bar"
+
+        def find_issues(self):
+            pass
+
+    Bar = register(Bar, task="regression")
+
+    assert REGISTRY["regression"].get("bar") == Bar