Fix pipeline unit tests

argilla-io · Oct 4, 2024 · 95f7618 · 95f7618
1 parent 09b588a
commit 95f7618
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 13 deletions.
diff --git a/src/distilabel/mixins/signature.py b/src/distilabel/mixins/signature.py
@@ -40,7 +40,7 @@ class SignatureMixin(BaseModel):
     """
 
     exclude_from_signature: Set[str] = Field(
-        _EXCLUDE_FROM_SIGNATURE_DEFAULTS, exclude=True
+        default=_EXCLUDE_FROM_SIGNATURE_DEFAULTS, exclude=True
     )
 
     @property

diff --git a/tests/unit/pipeline/test_base.py b/tests/unit/pipeline/test_base.py
@@ -145,12 +145,18 @@ def test_load_batch_manager(self, use_cache: bool) -> None:
 
             if use_cache:
                 mock_load_from_cache.assert_called_once_with(
-                    pipeline._cache_location["batch_manager"]
+                    dag=pipeline.dag,
+                    batch_manager_path=pipeline._cache_location["batch_manager"],
+                    steps_data_path=pipeline._cache_location["steps_data"],
                 )
                 mock_from_dag.assert_not_called()
             else:
                 mock_load_from_cache.assert_not_called()
-                mock_from_dag.assert_called_once_with(pipeline.dag)
+                mock_from_dag.assert_called_once_with(
+                    dag=pipeline.dag,
+                    use_cache=use_cache,
+                    steps_data_path=pipeline._cache_location["steps_data"],
+                )
 
     def test_setup_write_buffer(self) -> None:
         pipeline = DummyPipeline(name="unit-test-pipeline")
@@ -350,6 +356,7 @@ def test_run_stage_steps_and_wait(self, caplog) -> None:
 
             generator >> [step, step2] >> step3 >> step4
 
+        pipeline._load_batch_manager()
         pipeline._steps_load_status = {  # type: ignore
             generator.name: 1,
             step.name: 1,
@@ -373,6 +380,7 @@ def test_run_stage_steps_and_wait_with_failing_step(self, caplog) -> None:
             generator >> [step, step2] >> step3 >> step4
 
         pipeline._init_steps_load_status()
+        pipeline._load_batch_manager()
         pipeline._steps_load_status[generator.name] = _STEP_LOAD_FAILED_CODE  # type: ignore
         caplog.set_level(logging.INFO)
 
@@ -390,6 +398,7 @@ def test_run_stage_steps_and_wait_stop_called(self) -> None:
             generator >> [step, step2] >> step3 >> step4
 
         pipeline._init_steps_load_status()
+        pipeline._load_batch_manager()
         pipeline._stop_called = True
 
         assert pipeline._run_stage_steps_and_wait(stage=0) is False
@@ -649,7 +658,7 @@ def test_register_batch(self) -> None:
         pipeline._register_batch(batch)
 
         pipeline._batch_manager.register_batch.assert_called_once_with(
-            batch, path=pipeline._cache_location["batch_manager"]
+            batch, steps_data_path=pipeline._cache_location["steps_data"]
         )
 
     def test_send_last_batch_flag_to_step(self) -> None:
@@ -768,7 +777,7 @@ def test_handle_batch_on_stop(self) -> None:
         pipeline._handle_batch_on_stop(batch)
 
         batch_manager_mock.register_batch.assert_called_once_with(
-            batch, path=pipeline._cache_location["batch_manager"]
+            batch, steps_data_path=pipeline._cache_location["steps_data"]
         )
         batch_manager_mock.add_batch.assert_has_calls(
             [

diff --git a/tests/unit/steps/test_base.py b/tests/unit/steps/test_base.py
@@ -420,14 +420,6 @@ def test_step_dump(self) -> None:
                 "replicas": 1,
                 "resources": None,
             },
-            "exclude_from_signature": {
-                "disable_cuda_device_placement",
-                "gpu_memory_utilization",
-                "input_batch_size",
-                "resources",
-                "type_info",
-                "exclude_from_signature",
-            },
             "runtime_parameters_info": [
                 {
                     "name": "resources",