Merge pull request #439 from RelevanceAI/development

v0.32.0

Author: jtwinrelevanceai

ai_transform/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.31.3"
+__version__ = "0.32.0"
 
 from ai_transform.timer import Timer
 

ai_transform/engine/abstract_engine.py

@@ -1,5 +1,4 @@
 import time
-import logging
 import warnings
 
 from json import JSONDecodeError
@@ -8,7 +7,7 @@
 
 from tqdm.auto import tqdm
 
-from ai_transform.logger import format_logging_info, ic
+from ai_transform.logger import ic
 from ai_transform.types import Filter
 from ai_transform.dataset.dataset import Dataset
 from ai_transform.operator.abstract_operator import AbstractOperator
@@ -107,8 +106,10 @@ def __init__(
             filters = []
         assert isinstance(filters, list), "Filters must be applied as a list of Dictionaries"
 
-        if not refresh:
-            filters += self._get_refresh_filter(select_fields, dataset)
+        self._refresh = refresh
+        self._after_id = after_id
+
+        filters += self._get_refresh_filter()
         filters += self._get_workflow_filter()
 
         self._filters = filters
@@ -118,9 +119,6 @@ def __init__(
         else:
             self._size = dataset.len(filters=filters) if self._limit_documents is None else self._limit_documents
 
-        self._refresh = refresh
-        self._after_id = after_id
-
         self._successful_documents = 0
         self._success_ratio = None
 
@@ -206,36 +204,36 @@ def _operate(self, mini_batch):
             self._successful_documents += len(mini_batch)
             return transformed_batch
 
-    def _get_refresh_filter(self, select_fields: List[str], dataset: Dataset):
+    def _get_refresh_filter(self):
         # initialize the refresh filter container
-        refresh_filters = {"filter_type": "or", "condition_value": []}
+        input_field_filters = {"filter_type": "or", "condition_value": []}
 
         # initialize where the filters are going
-        input_field_filters = []
         output_field_filters = {"filter_type": "or", "condition_value": []}
 
-        # We want documents where all select_fields exists
+        # We want documents where any of the select_fields exists
         # as these are needed for operator ...
-        for field in select_fields:
-            input_field_filters += dataset[field].exists()
-
-        # ... and where any of its output_fields dont exist
-        for operator in self.operators:
-            if operator.output_fields is not None:
-                for output_field in operator.output_fields:
-                    output_field_filters["condition_value"] += dataset[output_field].not_exists()
-
         # We construct this as:
         #
-        #   input_field1 and input_field2 and (not output_field1 or not output_field2)
+        #   (input_field1 or input_field2) and (not output_field1 or not output_field2)
         #
         # This use case here is for two input fields and two output fields
         # tho this extends to arbitrarily many.
-        refresh_filters["condition_value"] = input_field_filters
-        refresh_filters["condition_value"] += [output_field_filters]
+        for field in self._select_fields:
+            input_field_filters["condition_value"] += self.dataset[field].exists()
+
+        # ... and where any of its output_fields dont exist
+        if not self._refresh:
+            for operator in self.operators:
+                if operator.output_fields is not None:
+                    for output_field in operator.output_fields:
+                        output_field_filters["condition_value"] += self.dataset[output_field].not_exists()
+
+            return [input_field_filters, output_field_filters]
 
-        # Wrap in list at end
-        return [refresh_filters]
+        else:
+            # Wrap in list at end
+            return [input_field_filters]
 
     def _get_workflow_filter(self, field: str = "_id"):
         # Get the required workflow filter as an environment variable

examples/fail_example.py

@@ -40,7 +40,8 @@ def __init__(
 
     def transform(self, documents: List[Document]) -> List[Document]:
         try:
-            text = [document[self.text_field] for document in documents]
+            raise ValueError
+
         except:
             # pass
             raise UserFacingError(

tests/conftest.py

@@ -63,6 +63,55 @@ def full_dataset(test_client: Client) -> Dataset:
     test_client.delete_dataset(dataset_id)
 
 
+@pytest.fixture(scope="class")
+def partial_dataset(test_client: Client) -> Dataset:
+    salt = "".join(random.choices(string.ascii_lowercase, k=10))
+    dataset_id = f"_sample_dataset_{salt}"
+    dataset = test_client.Dataset(dataset_id, expire=True)
+    documents = mock_documents(1000)
+    fields = ["sample_1_label", "sample_2_label", "sample_3_label"]
+    for document in documents:
+        for field in random.sample(fields, k=random.randint(1, 3)):
+            document.pop(field)
+    dataset.insert_documents(documents)
+    yield dataset
+    test_client.delete_dataset(dataset_id)
+
+
+@pytest.fixture(scope="class")
+def simple_partial_dataset(test_client: Client) -> Dataset:
+    salt = "".join(random.choices(string.ascii_lowercase, k=10))
+    dataset_id = f"_sample_dataset_{salt}"
+    dataset = test_client.Dataset(dataset_id, expire=True)
+    documents = mock_documents(1000)
+    fields = ["sample_1_label"]
+    for document in documents:
+        if random.random() < 0.5:
+            document.pop(fields[0])
+    dataset.insert_documents(documents)
+    yield dataset
+    test_client.delete_dataset(dataset_id)
+
+
+@pytest.fixture(scope="class")
+def partial_dataset_with_outputs(test_client: Client) -> Dataset:
+    salt = "".join(random.choices(string.ascii_lowercase, k=10))
+    dataset_id = f"_sample_dataset_{salt}"
+    dataset = test_client.Dataset(dataset_id, expire=True)
+    documents = mock_documents(1000)
+    fields = ["sample_1_label", "sample_2_label", "sample_3_label"]
+    for document in documents:
+        for field in random.sample(fields, k=random.randint(1, 3)):
+            document.pop(field)
+    for document in documents:
+        for field in fields:
+            if document.get(field) and random.random() < 0.5:
+                document[field + "_output"] = document[field] + "_output"
+    dataset.insert_documents(documents)
+    yield dataset
+    test_client.delete_dataset(dataset_id)
+
+
 @pytest.fixture(scope="class")
 def mixed_dataset(test_client: Client) -> Dataset:
     salt = "".join(random.choices(string.ascii_lowercase, k=10))
@@ -150,6 +199,26 @@ def transform(self, documents: DocumentList) -> DocumentList:
     return ExampleOperator()
 
 
+@pytest.fixture(scope="function")
+def test_partial_operator() -> AbstractOperator:
+    class PartialOperator(AbstractOperator):
+        def __init__(self, fields):
+            super().__init__(input_fields=fields, output_fields=[field + "_output" for field in fields])
+
+        def transform(self, documents: DocumentList) -> DocumentList:
+            """
+            Main transform function
+            """
+            for input_field, output_field in zip(self.input_fields, self.output_fields):
+                for document in documents:
+                    if document.get(input_field):
+                        document[output_field] = document[input_field] + "_output"
+
+            return documents
+
+    return PartialOperator
+
+
 @pytest.fixture(scope="function")
 def test_paid_operator() -> AbstractOperator:
     class ExampleOperator(AbstractOperator):
@@ -243,7 +312,7 @@ def test_user_facing_error_workflow_token(test_client: Client) -> str:
         job_id=job_id,
         dataset_id=dataset_id,
         authorizationToken=test_client.credentials.token,
-        text_field="sample_1_description_not_in_dataset",
+        text_field="sample_1_description",
     )
     config_string = json.dumps(config)
     config_bytes = config_string.encode()

tests/core/test_engine/test_stable_engine.py

@@ -1,20 +1,80 @@
+import uuid
+
+from typing import Type
+
 from ai_transform.dataset.dataset import Dataset
 from ai_transform.engine.stable_engine import StableEngine
 from ai_transform.engine.small_batch_stable_engine import SmallBatchStableEngine
 
 from ai_transform.operator.abstract_operator import AbstractOperator
-from ai_transform.workflow.abstract_workflow import AbstractWorkflow
+from ai_transform.workflow.abstract_workflow import Workflow
+
+
+def _random_id():
+    return str(uuid.uuid4())
 
 
 class TestStableEngine:
     def test_stable_engine(self, full_dataset: Dataset, test_operator: AbstractOperator):
         engine = StableEngine(full_dataset, test_operator, worker_number=0)
-        workflow = AbstractWorkflow(name="workflow_test123", engine=engine, job_id="test_job123")
+        workflow = Workflow(name=_random_id(), engine=engine, job_id=_random_id())
         workflow.run()
         assert engine.success_ratio == 1
 
     def test_small_batch_stable_engine(self, full_dataset: Dataset, test_operator: AbstractOperator):
         engine = SmallBatchStableEngine(full_dataset, test_operator)
-        workflow = AbstractWorkflow(name="workflow_test123", engine=engine, job_id="test_job123")
+        workflow = Workflow(name=_random_id(), engine=engine, job_id=_random_id())
+        workflow.run()
+        assert engine.success_ratio == 1
+
+
+class TestStableEngineFilters:
+    _SELECTED_FIELDS = ["sample_1_label", "sample_2_label", "sample_3_label"]
+
+    def test_stable_engine_filters1(self, partial_dataset: Dataset, test_partial_operator: Type[AbstractOperator]):
+        prev_health = partial_dataset.health()
+        operator = test_partial_operator(self._SELECTED_FIELDS)
+
+        engine = StableEngine(partial_dataset, operator, select_fields=self._SELECTED_FIELDS)
+        workflow = Workflow(name=_random_id(), engine=engine, job_id=_random_id())
+        workflow.run()
+
+        post_health = partial_dataset.health()
+        for input_field, output_field in zip(operator.input_fields, operator.output_fields):
+            assert prev_health[input_field]["exists"] == post_health[output_field]["exists"]
+
+        assert engine.success_ratio == 1
+
+    def test_stable_engine_filters2(
+        self, partial_dataset_with_outputs: Dataset, test_partial_operator: Type[AbstractOperator]
+    ):
+        prev_health = partial_dataset_with_outputs.health()
+        operator = test_partial_operator(self._SELECTED_FIELDS)
+
+        engine = StableEngine(
+            partial_dataset_with_outputs, operator, select_fields=self._SELECTED_FIELDS, refresh=False
+        )
+        workflow = Workflow(name=_random_id(), engine=engine, job_id=_random_id())
+        workflow.run()
+
+        post_health = partial_dataset_with_outputs.health()
+        for input_field, output_field in zip(operator.input_fields, operator.output_fields):
+            assert prev_health[input_field]["exists"] == post_health[output_field]["exists"]
+
+        assert engine.success_ratio == 1
+
+    def test_stable_engine_filters3(
+        self, simple_partial_dataset: Dataset, test_partial_operator: Type[AbstractOperator]
+    ):
+        prev_health = simple_partial_dataset.health()
+        operator = test_partial_operator(["sample_1_label"])
+
+        engine = StableEngine(simple_partial_dataset, operator, select_fields=["sample_1_label"], refresh=False)
+        workflow = Workflow(name=_random_id(), engine=engine, job_id=_random_id())
         workflow.run()
+
+        post_health = simple_partial_dataset.health()
+        for input_field, output_field in zip(operator.input_fields, operator.output_fields):
+            assert prev_health[input_field]["exists"] == post_health[output_field]["exists"]
+
         assert engine.success_ratio == 1