refactoring - alignment with automl

Signed-off-by: Lukasz Cmielowski <lcmielow@redhat.com>
Assisted-by: Cursor

Author: LukaszCmielowski

components/data_processing/automl/tabular_data_loader/component.py

@@ -343,7 +343,6 @@ def load_data_in_batches(
                 f"Available columns: {list(sampled_dataframe.columns)}"
             )
 
-
         sampled_dataframe.replace([math.inf, -math.inf], float("nan"), inplace=True)
 
         n_before_dedup = len(sampled_dataframe)

components/data_processing/autorag/documents_discovery/component.py

@@ -192,6 +192,9 @@ def _list_all_objects(s3_client):
 
             logger.info("Documents descriptor written to %s", descriptor_path)
 
+    if component_status is not None:
+        component_status.metadata["display_name"] = "Documents Discovery Status"
+
 
 if __name__ == "__main__":
     from kfp.compiler import Compiler

components/data_processing/autorag/test_data_loader/component.py

@@ -163,6 +163,9 @@ def _make_s3_client(verify=True):
                 record_count = len(benchmark_data)
                 logger.info("No sampling applied; record count: %s.", record_count)
 
+    if component_status is not None:
+        component_status.metadata["display_name"] = "Test Data Loader Status"
+
 
 if __name__ == "__main__":
     from kfp.compiler import Compiler

components/data_processing/autorag/text_extraction/component.py

@@ -464,6 +464,9 @@ def raise_if_threshold_exceeded(error_details: list, total_docs: int, tolerance:
             )
             raise_if_threshold_exceeded(all_error_details, total_docs, error_tolerance)
 
+    if component_status is not None:
+        component_status.metadata["display_name"] = "Text Extraction Status"
+
 
 if __name__ == "__main__":
     from kfp.compiler import Compiler

components/training/automl/autogluon_models_training/component.py

@@ -254,7 +254,7 @@ def _coerce_positive_class(value: Optional[str]) -> str | int | None:
             "completed",
             top_n=top_n,
             selected_models=top_models,
-            steps=["feature_engineering", "model_training", "stacking", "model_evaluation"],
+            steps=["feature_engineering", "model_training", "stacking", "evaluation"],
         )
 
         model_config = {

components/training/automl/autogluon_timeseries_models_training/README.md

@@ -125,7 +125,10 @@ def example_pipeline(
 
 ### Component status artifact
 
-Writes ``component_status.json`` under ``component_status`` with ``component_id`` ``autogluon_timeseries_models_training`` and training stages (``load_data``, ``model_selection`` with steps ``feature_engineering``, ``model_training``, ``stacking``, ``model_evaluation``, ``refit_and_evaluate``). Artifact metadata display name: **Timeseries Models Training Status**.
+Writes ``component_status.json`` under ``component_status`` with ``component_id``
+``autogluon_timeseries_models_training`` and training stages (``load_data``, ``model_selection`` with
+steps ``feature_engineering``, ``model_training``, ``stacking``, ``evaluation``, ``refit_and_evaluate``).
+Artifact metadata display name: **Timeseries Models Training Status**.
 
 Inference notebooks are loaded from ``shared/notebook_templates/timeseries_notebook.ipynb`` at runtime (same shared package data as tabular training).
 

components/training/automl/autogluon_timeseries_models_training/component.py

@@ -223,7 +223,7 @@ def autogluon_timeseries_models_training(
             "completed",
             top_n=top_n,
             selected_models=top_models,
-            steps=["feature_engineering", "model_training", "stacking", "model_evaluation"],
+            steps=["feature_engineering", "model_training", "stacking", "evaluation"],
         )
         logger.info(
             "Timeseries selection done: top_%s=%s best_score_test=%s",

components/training/automl/shared/component_status.py

@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+import base64
 import json
 import logging
 from collections.abc import Iterator
@@ -32,31 +33,49 @@
 STATUS_FAILED = "failed"
 
 
+class ComponentStatusEncoder(json.JSONEncoder):
+    """JSON encoder for component status metadata (Path, bytes, set, datetime)."""
+
+    def default(self, obj: Any) -> Any:
+        """Convert non-serializable objects to JSON-compatible types."""
+        if isinstance(obj, datetime):
+            return obj.isoformat().replace("+00:00", "Z") if obj.tzinfo is not None else obj.isoformat()
+        if isinstance(obj, Path):
+            return str(obj)
+        if isinstance(obj, bytes):
+            return base64.b64encode(obj).decode("ascii")
+        if isinstance(obj, set):
+            return sorted(obj, key=str)
+        return super().default(obj)
+
+
+def utc_now_z() -> str:
+    """Return current UTC time as an ISO-8601 string with ``Z`` suffix."""
+    return datetime.now(UTC).isoformat().replace("+00:00", "Z")
+
+
 class ComponentStatusTracker:
     """Track stage-level progress within a single component.
 
     Publishes component-local status to an artifact without requiring workspace.
     Each component independently tracks its stages and metadata.
     """
 
-    def __init__(self, artifact_path: str, component_id: str) -> None:
+    def __init__(self, artifact_path: str | None, component_id: str) -> None:
         """Initialize the status tracker.
 
         Args:
             artifact_path: Path to the KFP artifact directory where status.json will be written.
+                When ``None``, tracking is disabled (e.g. unit tests without a mock artifact).
             component_id: Unique component identifier (e.g., "autogluon_models_training").
         """
-        self.artifact_path = Path(artifact_path)
+        self._enabled = artifact_path is not None
+        self.artifact_path = Path(artifact_path) if self._enabled else Path(".")
         self.component_id = component_id
         self.stages: list[dict[str, Any]] = []
-        self.started_at = self._utc_now_iso()
+        self.started_at = utc_now_z()
         self.metadata: dict[str, Any] = {}
 
-    @staticmethod
-    def _utc_now_iso() -> str:
-        """Return current UTC timestamp in ISO format."""
-        return datetime.now(UTC).isoformat().replace("+00:00", "Z")
-
     def record(self, stage_id: str, status: str, **metadata: Any) -> None:
         """Record or update a stage's status.
 
@@ -73,21 +92,18 @@ def record(self, stage_id: str, status: str, **metadata: Any) -> None:
                          steps=["feature_eng", "training", "stacking"],
                          models_trained=15)
         """
-        # Find existing stage or create new one
         existing_idx = next((i for i, s in enumerate(self.stages) if s["id"] == stage_id), None)
 
         stage_data = {
             "id": stage_id,
             "status": status,
-            "timestamp": self._utc_now_iso(),
+            "timestamp": utc_now_z(),
             **metadata,
         }
 
         if existing_idx is not None:
-            # Update existing stage, preserving previously recorded metadata
             self.stages[existing_idx].update(stage_data)
         else:
-            # Append new stage
             self.stages.append(stage_data)
 
         logger.info(
@@ -115,19 +131,22 @@ def save(self) -> None:
         Creates the artifact directory if needed and writes component_status.json
         with all recorded stages and metadata.
         """
+        if not self._enabled:
+            return
+
         self.artifact_path.mkdir(parents=True, exist_ok=True)
 
         data = {
             "component_id": self.component_id,
             "started_at": self.started_at,
-            "completed_at": self._utc_now_iso(),
+            "completed_at": utc_now_z(),
             "stages": self.stages,
             "metadata": self.metadata,
         }
 
         output_file = self.artifact_path / COMPONENT_STATUS_FILENAME
         with output_file.open("w", encoding="utf-8") as f:
-            json.dump(data, f, indent=2)
+            json.dump(data, f, indent=2, cls=ComponentStatusEncoder)
 
         logger.info(
             "COMPONENT_STATUS component=%s saved status with %d stages to %s",
@@ -146,31 +165,39 @@ def save_best_effort(self) -> None:
                 self.component_id,
             )
 
-    def mark_active_failed(self, error: str) -> None:
+    def mark_active_failed(self, error: str | BaseException) -> None:
         """Mark the in-progress stage as failed, or the last open stage if none is active."""
+        if isinstance(error, BaseException):
+            error_msg = f"{type(error).__name__}: {error}" if str(error) else type(error).__name__
+        else:
+            error_msg = error
+
         active_statuses = (STATUS_STARTED, STATUS_RUNNING)
         for stage in reversed(self.stages):
             if stage.get("status") in active_statuses:
-                self.record(stage["id"], STATUS_FAILED, error=error)
+                self.record(stage["id"], STATUS_FAILED, error=error_msg)
                 return
 
         if self.stages and self.stages[-1].get("status") != STATUS_COMPLETED:
-            self.record(self.stages[-1]["id"], STATUS_FAILED, error=error)
+            self.record(self.stages[-1]["id"], STATUS_FAILED, error=error_msg)
             return
 
-        self.set_metadata(status=STATUS_FAILED, error=error)
+        self.set_metadata(status=STATUS_FAILED, error=error_msg)
 
     @contextmanager
     def stage(self, stage_id: str, **start_metadata: Any) -> Iterator[None]:
         """Record stage started/completed, or failed when an exception escapes the block."""
         self.record(stage_id, STATUS_STARTED, **start_metadata)
         try:
             yield
-        except Exception as exc:
-            self.record(stage_id, STATUS_FAILED, error=str(exc))
+        except BaseException as exc:
+            error_msg = f"{type(exc).__name__}: {exc}" if str(exc) else type(exc).__name__
+            self.record(stage_id, STATUS_FAILED, error=error_msg)
             raise
         else:
-            self.record(stage_id, STATUS_COMPLETED)
+            latest = next((s for s in reversed(self.stages) if s["id"] == stage_id), None)
+            if latest is None or latest.get("status") not in (STATUS_COMPLETED, STATUS_FAILED):
+                self.record(stage_id, STATUS_COMPLETED)
 
     def __enter__(self) -> ComponentStatusTracker:
         """Enter context: return this tracker."""
@@ -179,7 +206,7 @@ def __enter__(self) -> ComponentStatusTracker:
     def __exit__(self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: Any) -> bool:
         """On exit, mark active stage failed and save status best-effort."""
         if exc is not None:
-            self.mark_active_failed(str(exc))
+            self.mark_active_failed(exc)
         self.save_best_effort()
         return False
 
@@ -194,7 +221,7 @@ def load_component_status(artifact_path: str) -> dict[str, Any]:
 
     Returns:
         Dict containing component_id, started_at, completed_at, stages, and metadata.
-        Returns empty dict if file doesn't exist.
+        Returns empty dict if file doesn't exist or is unreadable.
 
     Example:
         status = load_component_status("/path/to/artifact")
@@ -204,5 +231,9 @@ def load_component_status(artifact_path: str) -> dict[str, Any]:
     if not status_file.exists():
         return {}
 
-    with status_file.open("r", encoding="utf-8") as f:
-        return json.load(f)
+    try:
+        with status_file.open("r", encoding="utf-8") as f:
+            return json.load(f)
+    except (json.JSONDecodeError, OSError) as e:
+        logger.warning("Failed to load status from %s: %s", status_file, e)
+        return {}

components/training/automl/shared/run_status_templates/pipelines/autogluon-tabular-training-pipeline.json

@@ -31,7 +31,7 @@
             "feature_engineering",
             "model_training",
             "stacking",
-            "model_evaluation"
+            "evaluation"
           ]
         },
         {

components/training/automl/shared/run_status_templates/pipelines/autogluon-timeseries-training-pipeline.json

@@ -31,7 +31,7 @@
             "feature_engineering",
             "model_training",
             "stacking",
-            "model_evaluation"
+            "evaluation"
           ]
         },
         {