aggregate and align the stages

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

Author: LukaszCmielowski

components/data_processing/automl/tabular_data_loader/README.md

@@ -224,8 +224,8 @@ load_task = automl_data_loader(
 
 In the tabular training pipeline, this component writes ``component_status.json`` under the
 ``component_status`` output artifact. The file includes ``component_id`` (``automl_data_loader``),
-``started_at``, ``completed_at``, a ``stages`` list (ids such as ``validate_inputs``,
-``read_and_sample``, ``cleanse``, ``split``, ``write_outputs``), and optional ``metadata``.
+``started_at``, ``completed_at``, a ``stages`` list (ids such as ``prepare_data``,
+``split_and_export``), and optional ``metadata``.
 Match stage ids to the tabular pipeline entry in ``component_stage_map.json`` from the
 ``publish-component-stage-map`` task.
 

components/data_processing/automl/tabular_data_loader/component.py

@@ -132,10 +132,7 @@ def automl_data_loader(  # noqa: D417
     # Initialize status tracker
     status = ComponentStatusTracker(component_status.path, "automl_data_loader")
     with status:
-        # Stage: validate_inputs
-        status.record("validate_inputs", "started")
-        # Validation happens inline below
-        status.record("validate_inputs", "completed")
+        status.record("prepare_data", "started")
 
         if sampling_method is None:
             if task_type in ("binary", "multiclass"):
@@ -324,10 +321,9 @@ def load_data_in_batches(
                 return _sample_random(text_stream, PANDAS_CHUNK_SIZE, max_size_bytes)
             return _sample_first_n_rows(text_stream, PANDAS_CHUNK_SIZE, max_size_bytes)
 
-        # Stage: read_and_sample
         status.record(
-            "read_and_sample",
-            "started",
+            "prepare_data",
+            "running",
             sampling_method=sampling_method,
             source=f"s3://{bucket_name}/{file_key}",
         )
@@ -347,10 +343,6 @@ def load_data_in_batches(
                 f"Available columns: {list(sampled_dataframe.columns)}"
             )
 
-        status.record("read_and_sample", "completed", rows=len(sampled_dataframe))
-
-        # Stage: cleanse
-        status.record("cleanse", "started")
 
         sampled_dataframe.replace([math.inf, -math.inf], float("nan"), inplace=True)
 
@@ -400,15 +392,14 @@ def load_data_in_batches(
             sampling_method,
         )
         status.record(
-            "cleanse",
+            "prepare_data",
             "completed",
             rows=n_samples,
             duplicates_dropped=n_dup_dropped,
             labels_dropped=n_dropped,
         )
 
-        # Stage: split
-        status.record("split", "started")
+        status.record("split_and_export", "started")
 
         # --- Train/test split ---
         from pathlib import Path
@@ -465,17 +456,13 @@ def load_data_in_batches(
         X_y_test.to_csv(sampled_test_dataset.path, index=False)
 
         status.record(
-            "split",
+            "split_and_export",
             "completed",
             test_size=test_size,
             selection_train_size=selection_train_size,
             stratify=stratify_effective,
         )
 
-        # Stage: write_outputs
-        status.record("write_outputs", "started")
-        status.record("write_outputs", "completed")
-
         component_status.metadata["display_name"] = "Data Loader Status"
 
         # Sample row for downstream use (JSON string to avoid NaN issues)

components/data_processing/automl/tabular_data_loader/tests/test_component_unit.py

@@ -212,8 +212,8 @@ def test_writes_component_status_json(self, tmp_path, monkeypatch):
         payload = json.loads(status_path.read_text())
         assert payload["component_id"] == "automl_data_loader"
         stage_ids = [stage["id"] for stage in payload["stages"]]
-        assert "read_and_sample" in stage_ids
-        assert "split" in stage_ids
+        assert "prepare_data" in stage_ids
+        assert "split_and_export" in stage_ids
 
     @mock.patch.dict("os.environ", mocked_env_variables)
     def test_sets_component_status_display_name(self, tmp_path):

components/data_processing/automl/timeseries_data_loader/README.md

@@ -104,6 +104,6 @@ def example_pipeline(
 
 In the time series training pipeline, this component writes ``component_status.json`` under the
 ``component_status`` output artifact. The file includes ``component_id`` (``timeseries_data_loader``),
-timestamps, and per-stage status (e.g. ``validate_inputs``, ``read_and_sample``, ``split``,
-``write_outputs``). Dashboards align stage ids with ``component_stage_map.json`` from
+timestamps, and per-stage status (e.g. ``prepare_data``, ``split_and_export``).
+Dashboards align stage ids with ``component_stage_map.json`` from
 ``publish-component-stage-map``.

components/data_processing/automl/timeseries_data_loader/component.py

@@ -96,8 +96,7 @@ def timeseries_data_loader(
 
     status = ComponentStatusTracker(component_status.path, "timeseries_data_loader")
     with status:
-        status.record("validate_inputs", "started")
-        status.record("validate_inputs", "completed")
+        status.record("prepare_data", "started")
 
         def get_s3_client(verify=True):
             """Create and return an S3 client using credentials from environment variables."""
@@ -289,8 +288,8 @@ def _clean_timeseries_dataframe(data, id_col, ts_col, log):
             return out.reset_index(drop=True)
 
         status.record(
-            "read_and_sample",
-            "started",
+            "prepare_data",
+            "running",
             source=f"s3://{bucket_name}/{file_key}",
         )
         df = load_timeseries_data_truncate(bucket_name, file_key, MAX_SIZE_BYTES, PANDAS_CHUNK_SIZE)
@@ -308,9 +307,6 @@ def _clean_timeseries_dataframe(data, id_col, ts_col, log):
                 f"with columns {sorted(required_columns)}."
             )
 
-        status.record("read_and_sample", "completed", rows=len(df))
-        status.record("cleanse", "started")
-
         df = _clean_timeseries_dataframe(df, id_column, timestamp_column, logger)
 
         n_valid = len(df)
@@ -321,8 +317,8 @@ def _clean_timeseries_dataframe(data, id_col, ts_col, log):
                 "Provide a larger dataset or fix invalid timestamps, null ids, and duplicate keys."
             )
 
-        status.record("cleanse", "completed", rows=n_valid)
-        status.record("split", "started")
+        status.record("prepare_data", "completed", rows=n_valid)
+        status.record("split_and_export", "started")
 
         # Create workspace datasets directory
         datasets_dir = Path(workspace_path) / "datasets"
@@ -391,7 +387,7 @@ def _concat_sorted(parts: list, sort_by: list) -> pd.DataFrame:
             )
 
         status.record(
-            "split",
+            "split_and_export",
             "completed",
             test_size=test_size,
             selection_train_size=selection_train_size,
@@ -424,8 +420,6 @@ def _concat_sorted(parts: list, sort_by: list) -> pd.DataFrame:
             "selection_train_size": selection_train_size,
         }
 
-        status.record("write_outputs", "started")
-        status.record("write_outputs", "completed")
         component_status.metadata["display_name"] = "Timeseries Data Loader Status"
 
         # Sample row for downstream use (JSON string to avoid NaN issues)

components/data_processing/autorag/documents_discovery/component.py

@@ -86,7 +86,7 @@ def get_test_data_docs_names() -> list[str]:
     _spec.loader.exec_module(_status_module)
     status = _status_module.bootstrap_status_tracker(embedded_artifact, component_status, "documents_discovery")
     with status:
-        with status.stage("validate_inputs"):
+        with status.stage("discover_documents"):
             s3_creds = {k: os.environ.get(k) for k in ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_S3_ENDPOINT"]}
             for k, v in s3_creds.items():
                 if v is None:
@@ -105,7 +105,6 @@ def _make_s3_client(verify=True):
                     verify=verify,
                 )
 
-        with status.stage("list_and_sample"):
             # Use paginator to handle buckets with >1,000 objects
             def _list_all_objects(s3_client):
                 """List all objects under prefix using pagination."""
@@ -186,7 +185,6 @@ def _list_all_objects(s3_client):
                 f"enabled_max={sampling_max_size}GB" if sampling_enabled else "disabled",
             )
 
-        with status.stage("write_descriptor"):
             os.makedirs(discovered_documents.path, exist_ok=True)
             descriptor_path = os.path.join(discovered_documents.path, DOCUMENTS_DESCRIPTOR_FILENAME)
             with open(descriptor_path, "w") as f:

components/data_processing/autorag/test_data_loader/component.py

@@ -78,7 +78,7 @@ class TestDataLoaderException(Exception):
     _spec.loader.exec_module(_status_module)
     status = _status_module.bootstrap_status_tracker(embedded_artifact, component_status, "test_data_loader")
     with status:
-        with status.stage("validate_inputs"):
+        with status.stage("load_benchmark"):
             if not test_data_bucket_name:
                 raise TypeError("test_data_bucket_name must be a non-empty string")
 
@@ -93,17 +93,16 @@ class TestDataLoaderException(Exception):
 
             s3_creds["AWS_DEFAULT_REGION"] = os.environ.get("AWS_DEFAULT_REGION")
 
-        def _make_s3_client(verify=True):
-            return boto3.client(
-                "s3",
-                endpoint_url=s3_creds["AWS_S3_ENDPOINT"],
-                region_name=s3_creds["AWS_DEFAULT_REGION"],
-                aws_access_key_id=s3_creds["AWS_ACCESS_KEY_ID"],
-                aws_secret_access_key=s3_creds["AWS_SECRET_ACCESS_KEY"],
-                verify=verify,
-            )
-
-        with status.stage("download_and_sample"):
+            def _make_s3_client(verify=True):
+                return boto3.client(
+                    "s3",
+                    endpoint_url=s3_creds["AWS_S3_ENDPOINT"],
+                    region_name=s3_creds["AWS_DEFAULT_REGION"],
+                    aws_access_key_id=s3_creds["AWS_ACCESS_KEY_ID"],
+                    aws_secret_access_key=s3_creds["AWS_SECRET_ACCESS_KEY"],
+                    verify=verify,
+                )
+
             s3_client = _make_s3_client()
 
             logger.info("Fetching test data from S3: bucket='%s', path='%s'.", test_data_bucket_name, test_data_path)
@@ -149,7 +148,6 @@ def _make_s3_client(verify=True):
                         f"Make sure that each test data records contains following keys: {benchmark_record_keys}."
                     )
 
-        with status.stage("write_output"):
             if 0 < benchmark_sample_size < len(benchmark_data) and isinstance(benchmark_data, list):
                 import random
 

components/data_processing/autorag/text_extraction/component.py

@@ -359,7 +359,7 @@ def raise_if_threshold_exceeded(error_details: list, total_docs: int, tolerance:
     with status:
         descriptor_path = Path(documents_descriptor.path) / DOCUMENTS_DESCRIPTOR_FILENAME
 
-        with status.stage("load_descriptor"):
+        with status.stage("extract_documents"):
             if not descriptor_path.exists():
                 raise FileNotFoundError(f"documents_descriptor.json not found at {descriptor_path}")
 
@@ -390,7 +390,6 @@ def raise_if_threshold_exceeded(error_details: list, total_docs: int, tolerance:
                 logger.info("No documents to process.")
                 return
 
-        with status.stage("extract_documents"):
             documents = sorted(documents, key=lambda d: d.get("size_bytes", 0), reverse=True)
 
             if max_extraction_workers is not None:

components/training/automl/autogluon_models_training/README.md

@@ -66,7 +66,7 @@ mutates predictor state. All artifacts are written under a single output artifac
 
 Writes ``component_status.json`` under the ``component_status`` output artifact with ``component_id``
 ``autogluon_models_training`` and stages such as ``load_data``, ``model_selection`` (optional ``steps``
-when completed), ``refit_full``, and ``evaluate_models``. Artifact metadata display name:
+when completed), ``refit_and_evaluate``. Artifact metadata display name:
 **Models Training Status**.
 
 ## Usage Examples 💡

components/training/automl/autogluon_models_training/component.py

@@ -299,9 +299,8 @@ def retrieve_pipeline_name(name: str) -> str:
         predictor_clone = predictor.clone(path=work_path, return_clone=True, dirs_exist_ok=True)
 
         # Refit all top models in a single call:  AutoGluon resolves stacking dependencies internally.
-        status.record("refit_full", "started")
+        status.record("refit_and_evaluate", "started")
         predictor_clone.refit_full(model=top_models, train_data_extra=extra_train_df)
-        status.record("refit_full", "completed", model_count=len(model_names_full))
 
         def replace_placeholder_in_notebook(notebook, replacements):
             for cell in notebook.get("cells", []):
@@ -657,7 +656,12 @@ def _process_model(model_name_full: str) -> tuple[str, dict]:
             "models": models_metadata,
         }
 
-        status.record("evaluate_models", "completed", eval_metric=str(predictor.eval_metric))
+        status.record(
+            "refit_and_evaluate",
+            "completed",
+            model_count=len(model_names_full),
+            eval_metric=str(predictor.eval_metric),
+        )
         component_status.metadata["display_name"] = "Models Training Status"
 
     return NamedTuple("outputs", eval_metric=str)(eval_metric=eval_metric)