Repository: caikit/caikit (python)

Issue:
Training endpoints change with multitask module
## Describe the bug

In an attempt to fix [a bug on the caikit-nlp side](https://github.com/caikit/caikit-nlp/issues/350) where tasks added to modules did not seem invokable, an unintentional change was triggered for the training side. The multiple tasks were relevant for inference, but the training endpoint also changed (from text generation to tokenization task).

## Expected behavior

Training endpoints should not be unexpectedly changed with task additions.

## Observed behavior

Prompt tuning changed from `TextGenerationTaskPeftPromptTuningTrain` to `TokenizationTaskPeftPromptTuningTrain`



Pull request description:
Multitask train modules 707
**What this PR does / why we need it**:

Closes #707 

This PR updates the logic for how the `ModuleClass.tasks` property is implemented. The previous use of a `set` was there to ensure uniqueness. It also had the subtle implication that all tasks are equally important to a given `ModuleClass` since a `set` does not guarantee order.

This second implication, however, was not respected in the train API generation. Since a given `ModuleClass` can have only a single `train` method, the mapping from `ModuleClass` to `task` was nondeterministic for multitask module training. With this PR, the order of `tasks` is now strictly defined to be the unique concatenation of the list of tasks in the module itself with the task lists of all parent modules in MRO.

**Special notes for your reviewer**:

**If applicable**:
- [ ] this PR contains documentation
- [x] this PR contains unit tests
- [x] this PR has been tested for backwards compatibility


Patch:
diff --git a/caikit/core/modules/decorator.py b/caikit/core/modules/decorator.py
index ca2d5007..458c13e6 100644
--- a/caikit/core/modules/decorator.py
+++ b/caikit/core/modules/decorator.py
@@ -188,15 +188,14 @@ def module(
         cls_.MODULE_CLASS = classname
         cls_.PRODUCER_ID = ProducerId(cls_.MODULE_NAME, cls_.MODULE_VERSION)
 
-        cls_._TASK_CLASSES = tasks
-
         # Parse the `train` and `run` signatures
         cls_.RUN_SIGNATURE = CaikitMethodSignature(cls_, "run")
         cls_.TRAIN_SIGNATURE = CaikitMethodSignature(cls_, "train")
         cls_._TASK_INFERENCE_SIGNATURES = {}
 
         # If the module has tasks, validate them:
-        for t in cls_._TASK_CLASSES:
+        task_classes = tasks
+        for t in task_classes:
             if not t.has_inference_method_decorators(module_class=cls_):
                 # Hackity hack hack - make sure at least one flavor is supported
                 validated = False
@@ -231,7 +230,16 @@ def module(
                 tasks_in_hierarchy.extend(class_._TASK_CLASSES)
 
         if tasks_in_hierarchy:
-            cls_._TASK_CLASSES += tasks_in_hierarchy
+            task_classes += tasks_in_hierarchy
+
+        # Make sure the tasks are unique. Note that the order here is important
+        # so that iterating the list of tasks is deterministic, unique, and the
+        # tasks given in the class' module list are shown before tasks inherited
+        # from parent classes.
+        cls_._TASK_CLASSES = []
+        for task in task_classes:
+            if task not in cls_._TASK_CLASSES:
+                cls_._TASK_CLASSES.append(task)
 
         # If no backend support described in the class, add current backend
         # as the only backend that can load models trained by this module
diff --git a/caikit/core/modules/meta.py b/caikit/core/modules/meta.py
index 424ed721..9b22795e 100644
--- a/caikit/core/modules/meta.py
+++ b/caikit/core/modules/meta.py
@@ -66,7 +66,7 @@ is known.
 """
 
 # Standard
-from typing import TYPE_CHECKING, Set
+from typing import TYPE_CHECKING, List
 import abc
 import functools
 
@@ -158,8 +158,8 @@ class _ModuleBaseMeta(abc.ABCMeta):
         return super().__new__(mcs, name, bases, attrs)
 
     @property
-    def tasks(cls) -> Set["TaskBase"]:
-        return set(cls._TASK_CLASSES)
+    def tasks(cls) -> List["TaskBase"]:
+        return [task for task in cls._TASK_CLASSES]
 
     def __setattr__(cls, name, val):
         """Overwrite __setattr__ to warn on any dynamic updates to the load function.


Test patch:
diff --git a/tests/core/test_task.py b/tests/core/test_task.py
index 0e509bd8..a6afaa3b 100644
--- a/tests/core/test_task.py
+++ b/tests/core/test_task.py
@@ -8,8 +8,14 @@ import pytest
 
 # Local
 from caikit.core import TaskBase, task
+from caikit.interfaces.common.data_model import File
 from sample_lib import SampleModule
-from sample_lib.data_model.sample import SampleInputType, SampleOutputType, SampleTask
+from sample_lib.data_model.sample import (
+    OtherOutputType,
+    SampleInputType,
+    SampleOutputType,
+    SampleTask,
+)
 from sample_lib.modules.multi_task import FirstTask, MultiTaskModule, SecondTask
 import caikit.core
 
@@ -171,7 +177,7 @@ def test_task_is_not_required_for_modules():
     class Stuff(caikit.core.ModuleBase):
         pass
 
-    assert Stuff.tasks == set()
+    assert Stuff.tasks == []
 
 
 def test_raises_if_tasks_not_list():
@@ -611,6 +617,32 @@ def test_validation_does_not_allow_union_supersets():
                 pass
 
 
+def test_tasks_property_order():
+    """Ensure that the tasks returned by .tasks have a deterministic order that
+    respects the order given in the module decorator
+    """
+    assert MultiTaskModule.tasks == [FirstTask, SecondTask]
+
+
+def test_tasks_property_unique():
+    """Ensure that entries in the tasks list is unique even when inherited from
+    modules with the same tasks
+    """
+
+    @caikit.core.module(
+        id=str(uuid.uuid4()),
+        name="DerivedMultitaskModule",
+        version="0.0.1",
+        task=SecondTask,
+    )
+    class DerivedMultitaskModule(MultiTaskModule):
+        @SecondTask.taskmethod()
+        def run_second_task(self, file_input: File) -> OtherOutputType:
+            return OtherOutputType("I'm a derivative!")
+
+    assert DerivedMultitaskModule.tasks == [SecondTask, FirstTask]
+
+
 # ----------- BACKWARDS COMPATIBILITY ------------------------------------------- ##