Embed MLflow run ID in dashboard CSV and Slack notifications - #2

.gitignore

@@ -162,4 +162,7 @@ Thumbs.db
 temp/
 
 # FORGE launcher configuration (personal settings)
-projects/core/launcher/launcher_config.yaml
+projects/core/launcher/launcher_config.yaml
+
+fournos-job-*.yaml
+kubeconfig*

projects/caliper/engine/file_export/artifacts_export_run.py

@@ -324,6 +324,7 @@ def run_multi_run_artifacts_export(
     run_dirs: list[Path],
     backend: tuple[str, ...] | list[str],
     mlflow_experiment: str | None = None,
+    mlflow_run_id: str | None = None,
     mlflow_run_name: str | None = None,
     mlflow_secrets_path: Path | None = None,
     mlflow_config_data: dict[str, Any] | None = None,
@@ -411,6 +412,7 @@ def run_multi_run_artifacts_export(
             parameters_file="parameters.json",
             tracking_uri=tracking_uri,
             experiment=experiment,
+            run_id=mlflow_run_id,
             parent_run_name=run_name,
             insecure_tls=insecure_tls,
             connection=mlflow_connection,

projects/caliper/engine/file_export/mlflow_backend.py

@@ -416,13 +416,15 @@ def _run(uri: str | None) -> tuple[str, dict[str, Any] | None]:
         start_kw: dict[str, Any] = {}
         if run_id:
             start_kw["run_id"] = run_id
-        elif run_name:
+        if run_name:
             start_kw["run_name"] = run_name
 
         meta: dict[str, Any] | None = None
         client = mlflow.tracking.MlflowClient()
         with mlflow.start_run(**start_kw):
             rid = mlflow.active_run().info.run_id
+            if run_id and run_name:
+                mlflow.set_tag("mlflow.runName", run_name)
             _apply_run_metadata(effective_meta)
             _apply_log_model(artifact_root, effective_meta, verbose=verbose)
 
@@ -491,6 +493,7 @@ def log_multi_run_artifacts(
     parameters_file: str,
     tracking_uri: str | None,
     experiment: str | None,
+    run_id: str | None = None,
     parent_run_name: str | None = None,
     insecure_tls: bool = False,
     connection: dict[str, Any] | None = None,
@@ -552,11 +555,15 @@ def _run(uri: str | None) -> tuple[str, dict[str, Any] | None]:
         client = mlflow.tracking.MlflowClient()
 
         start_kw: dict[str, Any] = {}
+        if run_id:
+            start_kw["run_id"] = run_id
         if parent_run_name:
             start_kw["run_name"] = parent_run_name
 
         with mlflow.start_run(**start_kw) as parent:
             parent_rid = parent.info.run_id
+            if run_id and parent_run_name:
+                mlflow.set_tag("mlflow.runName", parent_run_name)
             _apply_run_metadata(effective_meta)
 
             _upload_mlflow_files_parallel(

projects/caliper/orchestration/export.py

@@ -206,6 +206,9 @@ def run_from_orchestration_config(
 
     run_dirs = discover_run_dirs(from_path)
 
+    # Resume a pre-created MLflow run if the test step left a marker or test labels
+    mlflow_run_id = export_cfg.mlflow_run_id or _discover_precreated_mlflow_run_id(from_path)
+
     # Resolve descriptive run names from labels + run_naming config
     naming = resolve_run_names(
         run_dirs, mlflow_config_data, fallback_run_name=export_cfg.mlflow_run_name
@@ -226,6 +229,7 @@ def run_from_orchestration_config(
                 mlflow_run_name=naming.get("parent_run_name"),
                 mlflow_secrets_path=mlflow_secrets_path,
                 mlflow_config_data=mlflow_config_data,
+                mlflow_run_id=mlflow_run_id,
                 child_run_names=naming.get("child_run_names") or {},
                 verbose=export_cfg.verbose,
                 status_yaml_path=status_yaml,
@@ -238,7 +242,7 @@ def run_from_orchestration_config(
 
         mlflow_kwargs: dict[str, Any] = {
             "mlflow_experiment": export_cfg.mlflow_experiment,
-            "mlflow_run_id": export_cfg.mlflow_run_id,
+            "mlflow_run_id": mlflow_run_id,
             "mlflow_run_name": effective_name,
             "mlflow_secrets_path": mlflow_secrets_path,
         }
@@ -260,3 +264,161 @@ def run_from_orchestration_config(
 
     with open(status_yaml) as f:
         return yaml.safe_load(f.read())
+
+
+MLFLOW_PRECREATED_RUN_MARKER = "__mlflow_precreated_run__.yaml"
+"""Marker file written by the test step to communicate the pre-created MLflow
+run ID and experiment ID to the export step.
+
+The test step pre-creates an MLflow run (immediately ended) so that the
+``run_id`` and ``experiment_id`` are known *before* CSV generation.  These IDs
+are embedded in test labels and flow into the dashboard CSV.  The export step
+later discovers the marker via :func:`_discover_precreated_mlflow_run_id` and
+resumes the same run (``mlflow.start_run(run_id=...)``) to upload artifacts,
+avoiding a duplicate run.
+
+The marker is a YAML file with keys ``run_id`` and ``experiment_id``, written
+to ``ARTIFACT_DIR`` by :func:`precreate_mlflow_run`.
+"""
+
+
+def precreate_mlflow_run(
+    *,
+    secrets_path: Path,
+    experiment: str | None = None,
+    workspace: str | None = None,
+) -> dict[str, str]:
+    """Pre-create an MLflow run and write the marker file to ``ARTIFACT_DIR``.
+
+    The run is created and immediately ended (status FINISHED).  The export step
+    will resume it via ``mlflow.start_run(run_id=...)`` to upload artifacts.
+
+    The caller (test harness) is responsible for reading config and vault paths;
+    this function does not access the project config directly.
+
+    Returns a dict with ``run_id`` and ``experiment_id``.
+    """
+    import mlflow
+
+    from projects.caliper.public.file_export import (
+        load_mlflow_secrets_yaml,
+        mlflow_connection_env,
+    )
+
+    secrets_data = load_mlflow_secrets_yaml(secrets_path)
+    tracking_uri = secrets_data.get("tracking_uri", "")
+
+    prev_workspace = os.environ.get("MLFLOW_WORKSPACE")
+    prev_tracking_uri = mlflow.get_tracking_uri()
+    try:
+        with mlflow_connection_env(secrets_data):
+            if tracking_uri:
+                mlflow.set_tracking_uri(tracking_uri)
+            if workspace:
+                os.environ["MLFLOW_WORKSPACE"] = workspace
+            if experiment:
+                mlflow.set_experiment(experiment)
+
+            run_name = os.environ.get("FJOB_NAME")
+            with mlflow.start_run(run_name=run_name):
+                active = mlflow.active_run()
+                run_id = active.info.run_id
+                experiment_id = str(active.info.experiment_id)
+    finally:
+        if prev_workspace is not None:
+            os.environ["MLFLOW_WORKSPACE"] = prev_workspace
+        else:
+            os.environ.pop("MLFLOW_WORKSPACE", None)
+        mlflow.set_tracking_uri(prev_tracking_uri)
+
+    meta = {"run_id": run_id, "experiment_id": experiment_id}
+
+    marker_path = env.ARTIFACT_DIR / MLFLOW_PRECREATED_RUN_MARKER
+    marker_path.parent.mkdir(parents=True, exist_ok=True)
+    marker_path.write_text(yaml.safe_dump(meta, sort_keys=False), encoding="utf-8")
+    logger.info(
+        "Pre-created MLflow run %s (experiment=%s), marker: %s",
+        run_id,
+        experiment_id,
+        marker_path,
+    )
+
+    return meta
+
+
+def _read_mlflow_ids_from_marker() -> tuple[str, str]:
+    """Read run_id and experiment_id from the pre-created marker file on disk."""
+    artifact_dir = Path(env.ARTIFACT_DIR) if env.ARTIFACT_DIR else None
+    if not artifact_dir:
+        logger.warning("ARTIFACT_DIR not set, cannot read MLflow marker")
+        return "", ""
+    for marker in sorted(artifact_dir.rglob(MLFLOW_PRECREATED_RUN_MARKER)):
+        data = yaml.safe_load(marker.read_text(encoding="utf-8"))
+        if not isinstance(data, dict):
+            logger.warning(
+                "MLflow marker %s has unexpected type %s, skipping", marker, type(data).__name__
+            )
+            continue
+        if data.get("run_id"):
+            return data["run_id"], data.get("experiment_id", "")
+    logger.warning("No valid MLflow marker found under %s", artifact_dir)
+    return "", ""
+
+
+def build_mlflow_run_url(
+    *,
+    secrets_path: Path,
+    workspace: str | None = None,
+) -> str:
+    """Construct the MLflow run URL from vault secrets and the marker file.
+
+    The caller (test harness) is responsible for resolving ``secrets_path``
+    and ``workspace``; this function does not access the project config.
+    """
+    from urllib.parse import quote
+
+    from projects.caliper.public.file_export import (
+        assert_tracking_uri_has_no_userinfo,
+        load_mlflow_secrets_yaml,
+    )
+
+    run_id, experiment_id = _read_mlflow_ids_from_marker()
+    if not run_id or not experiment_id:
+        logger.warning("Cannot build MLflow URL: run_id or experiment_id missing from marker")
+        return ""
+
+    if not secrets_path.exists():
+        logger.warning("Cannot build MLflow URL: secrets file %s not found", secrets_path)
+        return ""
+
+    secrets_data = load_mlflow_secrets_yaml(secrets_path)
+    tracking_uri = secrets_data.get("tracking_uri", "").rstrip("/")
+    if not tracking_uri.startswith(("http://", "https://")):
+        logger.warning("Cannot build MLflow URL: tracking_uri has unsupported scheme")
+        return ""
+    assert_tracking_uri_has_no_userinfo(tracking_uri)
+
+    qs = f"?workspace={quote(workspace, safe='')}" if workspace else ""
+    return f"{tracking_uri}/#/experiments/{experiment_id}/runs/{run_id}/artifacts{qs}"
+
+
+def _discover_precreated_mlflow_run_id(from_path: Path) -> str | None:
+    """Find a pre-created MLflow run_id from the marker file."""
+    for marker in sorted(from_path.rglob(MLFLOW_PRECREATED_RUN_MARKER)):
+        try:
+            data = yaml.safe_load(marker.read_text(encoding="utf-8"))
+            if not isinstance(data, dict):
+                logger.warning(
+                    "MLflow marker %s has unexpected type %s, skipping",
+                    marker,
+                    type(data).__name__,
+                )
+                continue
+            run_id = data.get("run_id")
+            if run_id:
+                logger.info("Found pre-created MLflow run_id: %s (from marker %s)", run_id, marker)
+                return run_id
+        except (OSError, yaml.YAMLError) as e:
+            logger.warning("Failed to read MLflow pre-created run marker %s: %s", marker, e)
+
+    return None

projects/caliper/public/__init__.py

@@ -0,0 +1 @@
+"""Public API surface for caliper — safe for orchestration-layer imports."""

projects/caliper/public/file_export.py

@@ -0,0 +1,17 @@
+"""Public re-exports of caliper engine file_export utilities.
+
+Orchestration code should import from here instead of reaching into
+``projects.caliper.engine.file_export`` directly.
+"""
+
+from projects.caliper.engine.file_export.mlflow_secrets import (
+    assert_tracking_uri_has_no_userinfo,
+    load_mlflow_secrets_yaml,
+    mlflow_connection_env,
+)
+
+__all__ = [
+    "assert_tracking_uri_has_no_userinfo",
+    "load_mlflow_secrets_yaml",
+    "mlflow_connection_env",
+]

projects/rhaiis/orchestration/analysis.py

@@ -165,6 +165,11 @@ def run_regression_check(
             restrict_profiles=restrict_profiles,
         )
 
+        _ea = engine_args or {}
+        tp = _ea.get("tensor-parallel-size") or _ea.get("tp-size") or _ea.get("tp_size") or ""
+        dp = _ea.get("data-parallel-size") or _ea.get("dp-size") or ""
+        slack_user = config.project.get_config("tests.rhaiis.slack_user", "")
+
         if analysis.get("regression_count", 0) > 0 or analysis.get("improvement_count", 0) > 0:
             report_url = ""
             agent_cfg = config.project.get_config("rhaiis.agent_analysis", {})
@@ -181,20 +186,32 @@ def run_regression_check(
 
             from projects.rhaiis.postprocess.regression import send_regression_notification
 
-            _ea = engine_args or {}
-            tp = _ea.get("tensor-parallel-size") or _ea.get("tp-size") or _ea.get("tp_size") or ""
-            dp = _ea.get("data-parallel-size") or _ea.get("dp-size") or ""
             send_regression_notification(
                 analysis,
                 model=model_cfg.get("hf_model_id", ""),
                 accelerator=accelerator,
                 job_id=run_uuid,
-                slack_user=config.project.get_config("tests.rhaiis.slack_user", ""),
+                slack_user=slack_user,
                 notification_vault="psap-forge-notifications",
                 report_url=report_url,
                 tp=str(tp),
                 dp=str(dp),
             )
+        elif config.project.get_config("tests.rhaiis.slack_notify_always", False):
+            from projects.rhaiis.postprocess.regression import send_success_notification
+
+            send_success_notification(
+                model=model_cfg.get("hf_model_id", ""),
+                accelerator=accelerator,
+                job_id=run_uuid,
+                slack_user=slack_user,
+                notification_vault="psap-forge-notifications",
+                tp=str(tp),
+                dp=str(dp),
+                version=current_version,
+                workload_keys=config.project.get_config("tests.rhaiis.workload_keys", []),
+                cluster=config.project.get_config("rhaiis.cluster_tag", ""),
+            )
     except Exception:
         logger.warning("Regression analysis failed; continuing", exc_info=True)
     finally:

projects/rhaiis/orchestration/config.yaml

@@ -28,6 +28,7 @@ tests:
     version: ""
     compare_version: ""
     slack_user: ""
+    slack_notify_always: false
 
 caliper:
   postprocess: