refactor the code

src/syngen/VERSION

@@ -1 +1 @@
-0.9.52rc31
+0.9.52rc32

src/syngen/ml/worker/worker.py

@@ -300,11 +300,11 @@ def _train_table(self, table, metadata, delta):
             batch_size=train_settings["batch_size"],
             loader=self.loader
         )
-        self._write_success_message(slugify(table))
+        self._write_success_file(slugify(table))
         self._save_metadata_file()
         ProgressBarHandler().set_progress(
             delta=delta,
-            message=f"Training of the table - {table} was completed"
+            message=f"Training of the table - '{table}' was completed"
         )
 
     def __train_tables(
@@ -361,6 +361,35 @@ def _find_parent_table(self, table):
             ), None
         )
 
+    import os
+    from slugify import slugify
+
+    def _check_completion_of_training(self, table_name: str):
+        """
+        Check if the training process of a specific table has been completed.
+
+        Args:
+        table_name (str): The name of the table to check.
+
+        Raises:
+        FileNotFoundError: If the success file does not exist.
+        ValueError: If the content of the success file does not indicate success.
+        """
+        path_to_success_file = f"model_artifacts/resources/{slugify(table_name)}/message.success"
+        error_message = (
+            f"The training of the table - '{table_name}' hasn't been completed. "
+            "Please, retrain the table."
+        )
+
+        if not os.path.exists(path_to_success_file):
+            raise FileNotFoundError(error_message)
+
+        with open(path_to_success_file, 'r') as file:
+            content = file.read().strip()
+
+        if content != "SUCCESS":
+            raise ValueError(error_message)
+
     def _infer_table(self, table, metadata, type_of_process, delta, is_nested=False):
         """
         Infer process for a single table
@@ -395,7 +424,7 @@ def _infer_table(self, table, metadata, type_of_process, delta, is_nested=False)
         )
         ProgressBarHandler().set_progress(
             delta=delta,
-            message=f"Infer process of the table - {table} was completed"
+            message=f"Infer process of the table - '{table}' was completed"
         )
         MlflowTracker().end_run()
 
@@ -414,6 +443,7 @@ def __infer_tables(
 
         non_surrogate_tables = [table for table in tables if table not in self.divided]
         for table in non_surrogate_tables:
+            self._check_completion_of_training(slugify(table))
             self._infer_table(
                 table=table,
                 metadata=config_of_tables,
@@ -455,7 +485,7 @@ def _generate_reports(self):
         Report().clear_report()
 
     @staticmethod
-    def _write_success_message(table_name: str):
+    def _write_success_file(table_name: str):
         """
         Write success message to the '.success' file
         """

src/tests/conftest.py

@@ -154,6 +154,20 @@ def test_metadata_storage():
         shutil.rmtree("model_artifacts")
 
 
+@pytest.fixture
+def test_success_file():
+    path_to_test_dir = "model_artifacts/resources/test-table"
+    os.makedirs(path_to_test_dir, exist_ok=True)
+    success_file_path = f"{path_to_test_dir}/message.success"
+    with open(success_file_path, "w") as f:
+        f.write("PROGRESS")
+
+    yield success_file_path
+    if os.path.exists(success_file_path):
+        shutil.rmtree("model_artifacts")
+
+
+
 @pytest.fixture
 def test_metadata_file():
     return {

src/tests/unit/test_worker/test_worker.py

@@ -1,5 +1,7 @@
 from unittest.mock import patch, MagicMock
 
+import pytest
+
 from syngen.ml.worker import Worker
 from syngen.ml.config import Validator
 
@@ -1823,3 +1825,166 @@ def test_launch_train_with_metadata_without_train_settings(
         True
     )
     rp_logger.info(SUCCESSFUL_MESSAGE)
+
+
+
+@patch.object(Worker, "_generate_reports")
+@patch.object(Worker, "_check_completion_of_training", return_value=None)
+@patch.object(Worker, "_infer_table")
+@patch.object(Worker, "_collect_metrics_in_infer")
+@patch.object(Validator, "_validate_metadata")
+@patch.object(Validator, "_check_existence_of_destination")
+def test_launch_infer_pretrained_table(
+    mock_check_existence_of_destination,
+    mock_validate_metadata,
+    mock_collect_metrics_in_infer,
+    mock_infer_table,
+    mock_check_completion_of_training,
+    mock_generate_reports,
+    rp_logger,
+):
+    """
+    Test that the inference process has been launched
+    if the training process of the table has been finished
+    """
+    rp_logger.info(
+        "Test that the inference process has been launched "
+        "if the training process of the table has been finished"
+    )
+    worker = Worker(
+        table_name="test_table",
+        metadata_path=None,
+        settings={
+            "size": 300,
+            "run_parallel": True,
+            "random_seed": 3,
+            "reports": ["accuracy"],
+            "batch_size": 300,
+        },
+        log_level="INFO",
+        type_of_process="infer",
+        loader=None
+    )
+    worker.launch_infer()
+    mock_check_existence_of_destination.assert_called_once()
+    mock_validate_metadata.assert_called_once_with("test_table")
+    mock_infer_table.assert_called_once_with(
+        table="test_table",
+        metadata={
+            "test_table": {
+                "train_settings": {
+                    "source": None
+                },
+                'infer_settings': {
+                    "size": 300,
+                    "run_parallel": True,
+                    "random_seed": 3,
+                    "reports": ["accuracy"],
+                    "batch_size": 300
+                },
+                "keys": {},
+                "format": {}
+            }
+        },
+        type_of_process="infer",
+        delta=0.25
+    )
+    mock_collect_metrics_in_infer.assert_called_once_with(["test_table"])
+    mock_generate_reports.assert_called_once()
+
+
+@patch.object(Worker, "_infer_table")
+@patch.object(Validator, "_validate_metadata")
+@patch.object(Validator, "_check_existence_of_destination")
+def test_launch_infer_not_pretrained_table(
+    mock_check_existence_of_destination,
+    mock_validate_metadata,
+    mock_collect_metrics_in_infer,
+    mock_infer_table,
+    caplog,
+    rp_logger,
+):
+    """
+    Test that the inference process hasn't been started
+    in case the training process of the table hasn't been finished,
+    and the appropriate success file 'message.success' is absent
+    """
+    rp_logger.info(
+        "Test that the inference process hasn't been started "
+        "in case the training process of the table hasn't been finished, " 
+        "and the appropriate success file 'message.success' is absent"
+    )
+    worker = Worker(
+        table_name="test_table",
+        metadata_path=None,
+        settings={
+            "size": 300,
+            "run_parallel": True,
+            "random_seed": 3,
+            "reports": ["accuracy"],
+            "batch_size": 300,
+        },
+        log_level="INFO",
+        type_of_process="infer",
+        loader=None
+    )
+    with pytest.raises(FileNotFoundError):
+        with caplog.at_level("ERROR"):
+            worker.launch_infer()
+            assert (
+                "The training of the table - 'test-table' hasn't been completed. "
+                "Please, retrain the table."
+                in caplog.text
+            )
+    mock_infer_table.assert_not_called()
+    rp_logger.info(SUCCESSFUL_MESSAGE)
+
+
+@patch.object(Worker, "_infer_table")
+@patch.object(Validator, "_validate_metadata")
+@patch.object(Validator, "_check_existence_of_destination")
+def test_launch_infer_not_pretrained_table(
+    mock_check_existence_of_destination,
+    mock_validate_metadata,
+    mock_infer_table,
+    test_success_file,
+    caplog,
+    rp_logger,
+):
+    """
+    Test that the inference process hasn't been started
+    in case the training process of the table hasn't been finished,
+    and the appropriate success file 'message.success' is present,
+    but the content of the file doesn't correspond to finished training process
+    """
+    rp_logger.info(
+        "Test that the inference process hasn't been started "
+        "in case the training process of the table hasn't been finished, " 
+        "and the appropriate success file 'message.success' is present, "
+        "but the content of the file doesn't correspond to finished training process"
+    )
+    worker = Worker(
+        table_name="test_table",
+        metadata_path=None,
+        settings={
+            "size": 300,
+            "run_parallel": True,
+            "random_seed": 3,
+            "reports": ["accuracy"],
+            "batch_size": 300,
+        },
+        log_level="INFO",
+        type_of_process="infer",
+        loader=None
+    )
+    with pytest.raises(ValueError):
+        with caplog.at_level("ERROR"):
+            worker.launch_infer()
+            assert (
+                "The training of the table - 'test-table' hasn't been completed. "
+                "Please, retrain the table."
+                in caplog.text
+            )
+    mock_infer_table.assert_not_called()
+    rp_logger.info(SUCCESSFUL_MESSAGE)
+