More useful return values from runways.

Author: superdosh

src/modelplane/runways/annotator.py

@@ -27,7 +27,12 @@
     is_debug_mode,
     setup_annotator_credentials,
 )
-from modelplane.runways.data import BaseInput, build_and_log_input
+from modelplane.runways.data import (
+    Artifact,
+    BaseInput,
+    RunArtifacts,
+    build_and_log_input,
+)
 
 KNOWN_ENSEMBLES: Dict[str, AnnotatorSet] = {}
 # try to load the private ensemble
@@ -55,7 +60,7 @@ def annotate(
     prompt_text_col=None,
     sut_uid_col=None,
     sut_response_col=None,
-) -> str:
+) -> RunArtifacts:
     """
     Run annotations and record measurements.
     """
@@ -138,7 +143,15 @@ def annotate(
                 / pipeline_runner.output_file_name,
                 dir=tmp,
             )
-        return run.info.run_id
+            artifacts = {
+                input_data.local_path().name: input_data.artifact,
+                pipeline_runner.output_file_name: Artifact(
+                    experiment_id=run.info.experiment_id,
+                    run_id=run.info.run_id,
+                    name=pipeline_runner.output_file_name,
+                ),
+            }
+        return RunArtifacts(run_id=run.info.run_id, artifacts=artifacts)
 
 
 def _get_annotator_settings(

src/modelplane/runways/data.py

@@ -1,3 +1,4 @@
+from dataclasses import dataclass
 import os
 import shutil
 from abc import ABC, abstractmethod
@@ -14,13 +15,39 @@
 )
 
 
+class Artifact:
+
+    def __init__(self, experiment_id: str, run_id: str, name: str):
+        self.name = name
+        tracking_uri = mlflow.get_tracking_uri()
+        self._mlflow_link = (
+            f"{tracking_uri}/#/experiments/{experiment_id}/runs/{run_id}"
+        )
+        self._download_link = f"{tracking_uri}/api/2.0/mlflow/artifacts/download?run_id={run_id}&artifact_path={name}"
+
+    @property
+    def mlflow_link(self) -> str:
+        return self._mlflow_link
+
+    @property
+    def download_link(self) -> str:
+        return self._download_link
+
+
+@dataclass
+class RunArtifacts:
+    run_id: str
+    artifacts: dict[str, Artifact | None]
+
+
 class BaseInput(ABC):
     """Base class for input datasets."""
 
     input_type: str
 
     def __init__(self):
         self.input_run_id = None
+        self._artifact = None
 
     def __init_subclass__(cls):
         super().__init_subclass__()
@@ -39,17 +66,15 @@ def log_artifact(self):
         local = self.local_path()
         mlflow.log_artifact(str(local))
         mlflow.set_tags(self.input_tags())
-        mlflow_tracking_uri = mlflow.get_tracking_uri()
-        self._mlflow_link = f"{mlflow_tracking_uri}/#/experiments/{current_run.info.experiment_id}/runs/{current_run.info.run_id}/artifacts/{local.name}"
-        self._download_link = f"{mlflow_tracking_uri}/api/2.0/mlflow-artifacts/artifacts/{local.name}?run_id={current_run.info.run_id}"
-
-    @property
-    def mlflow_link(self) -> str:
-        return self._mlflow_link
+        self._artifact = Artifact(
+            experiment_id=current_run.info.experiment_id,
+            run_id=current_run.info.run_id,
+            name=local.name,
+        )
 
     @property
-    def download_link(self) -> str:
-        return self._download_link
+    def artifact(self) -> Artifact | None:
+        return self._artifact
 
     @abstractmethod
     def local_path(self) -> Path:

src/modelplane/runways/responder.py

@@ -17,7 +17,12 @@
     is_debug_mode,
     setup_sut_credentials,
 )
-from modelplane.runways.data import BaseInput, build_and_log_input
+from modelplane.runways.data import (
+    Artifact,
+    BaseInput,
+    RunArtifacts,
+    build_and_log_input,
+)
 
 
 def respond(
@@ -30,7 +35,7 @@ def respond(
     num_workers: int = 1,
     prompt_uid_col=None,
     prompt_text_col=None,
-) -> str:
+) -> RunArtifacts:
     secrets = setup_sut_credentials(sut_id)
     sut = SUTS.make_instance(uid=sut_id, secrets=secrets)
     params = {"num_workers": num_workers}
@@ -68,4 +73,13 @@ def respond(
                 local_path=pipeline_runner.output_dir()
                 / pipeline_runner.output_file_name,
             )
-        return run.info.run_id
+            artifacts = {
+                input_data.local_path().name: input_data.artifact,
+                pipeline_runner.output_file_name: Artifact(
+                    experiment_id=run.info.experiment_id,
+                    run_id=run.info.run_id,
+                    name=pipeline_runner.output_file_name,
+                ),
+            }
+
+        return RunArtifacts(run_id=run.info.run_id, artifacts=artifacts)

src/modelplane/runways/scorer.py

@@ -18,7 +18,7 @@
     RUN_TYPE_TAG_NAME,
     get_experiment_id,
 )
-from modelplane.runways.data import BaseInput, build_and_log_input
+from modelplane.runways.data import BaseInput, RunArtifacts, build_and_log_input
 
 
 def score(
@@ -30,7 +30,7 @@ def score(
     sample_uid_col: str | None = None,
     annotator_uid_col: str | None = ANNOTATION_SCHEMA.annotator_uid,
     annotation_col: str | None = ANNOTATION_SCHEMA.annotation,
-):
+) -> RunArtifacts:
     """
     Score annotations against ground truth.
     Annotations are expected to be in JSON format with an "is_safe" field.
@@ -78,6 +78,11 @@ def score(
             )
             mlflow.log_metric("num_ground_truth_samples", len(ground_truth_data.df))
 
+            artifacts = {
+                annotation_input.local_path().name: annotation_input.artifact,
+                ground_truth_input.local_path().name: ground_truth_input.artifact,
+            }
+
         # Score each annotator in the annotation dataframe.
         for annotator in annotation_data.annotators:
             score = score_annotator(annotator, annotation_data, ground_truth_data)
@@ -92,7 +97,7 @@ def score(
                 else:
                     mlflow.log_metric(f"{annotator}_{metric}", score[metric])
 
-        return run.info.run_id
+        return RunArtifacts(run_id=run.info.run_id, artifacts=artifacts)
 
 
 def score_annotator(annotator: str, annotation_data, ground_truth_data):

tests/it/runways/test_e2e.py

@@ -22,22 +22,22 @@ def test_e2e():
     experiment = "test_experiment_" + time.strftime("%Y%m%d%H%M%S", time.localtime())
     num_workers = 1
 
-    run_id = check_responder(
+    run_artifacts = check_responder(
         sut_id=sut_id,
         prompts=prompts,
         experiment=experiment,
         disable_cache=True,
         num_workers=num_workers,
     )
-    run_id = check_annotator(
-        response_run_id=run_id,
+    run_artifacts = check_annotator(
+        response_run_id=run_artifacts.run_id,
         annotator_ids=[TEST_ANNOTATOR_ID],
         experiment=experiment,
         disable_cache=True,
         num_workers=num_workers,
     )
     check_scorer(
-        annotation_run_id=run_id,
+        annotation_run_id=run_artifacts.run_id,
         ground_truth=ground_truth,
         annotator_id=TEST_ANNOTATOR_ID,
         experiment=experiment,
@@ -51,7 +51,7 @@ def check_responder(
     disable_cache: bool,
     num_workers: int,
 ):
-    run_id = respond(
+    run_artifacts = respond(
         sut_id=sut_id,
         prompts=prompts,
         experiment=experiment,
@@ -62,10 +62,10 @@ def check_responder(
     # confirm experiment exists
     exp = mlflow.get_experiment_by_name(experiment)
     assert exp is not None
-    assert run_id is not None
+    assert run_artifacts.run_id is not None
 
     # validate params / tags logged
-    run = mlflow.get_run(run_id)
+    run = mlflow.get_run(run_artifacts.run_id)
     params = run.data.params
     tags = run.data.tags
     assert params.get("num_workers") == str(num_workers)
@@ -75,7 +75,7 @@ def check_responder(
     with tempfile.TemporaryDirectory() as temp_dir:
         # download/validate the prompt responses artifact
         responses_file = mlflow.artifacts.download_artifacts(
-            run_id=run_id,
+            run_id=run_artifacts.run_id,
             artifact_path=PROMPT_RESPONSE_ARTIFACT_NAME,
             dst_path=temp_dir,
         )
@@ -90,7 +90,7 @@ def check_responder(
                 assert (
                     yesno.lower() == expected
                 ), f"Unexpectedly got '{yesno} for prompt '{response['prompt_text']}'"
-    return run_id
+    return run_artifacts
 
 
 def check_annotator(
@@ -101,7 +101,7 @@ def check_annotator(
     num_workers: int,
 ):
     # run the annotator
-    run_id = annotate(
+    run_artifacts = annotate(
         response_run_id=response_run_id,
         annotator_ids=annotator_ids,
         experiment=experiment,
@@ -113,7 +113,7 @@ def check_annotator(
     assert exp is not None
 
     # validate params / tags / metrics logged
-    run = mlflow.get_run(run_id)
+    run = mlflow.get_run(run_artifacts.run_id)
     params = run.data.params
     tags = run.data.tags
     metrics = run.data.metrics
@@ -129,12 +129,12 @@ def check_annotator(
     ), "Expected total_safe to be 5"
 
     # confirm annotations.csv exists
-    artifacts = mlflow.artifacts.list_artifacts(run_id=run_id)
+    artifacts = mlflow.artifacts.list_artifacts(run_id=run_artifacts.run_id)
     assert any(
         artifact.path == "annotations.csv" for artifact in artifacts
     ), "Expected 'annotations.csv' artifact not found in run"
     # TODO: validate annotations.csv
-    return run_id
+    return run_artifacts
 
 
 def check_scorer(
@@ -143,13 +143,13 @@ def check_scorer(
     annotator_id: str,
     experiment: str,
 ):
-    run_id = score(annotation_run_id, experiment, ground_truth)
+    run_artifacts = score(annotation_run_id, experiment, ground_truth)
     # confirm experiment exists
     exp = mlflow.get_experiment_by_name(experiment)
     assert exp is not None
 
     # validate params / metrics logged
-    run = mlflow.get_run(run_id)
+    run = mlflow.get_run(run_artifacts.run_id)
     params = run.data.params
     metrics = run.data.metrics
     assert params.get("annotation_run_id") == annotation_run_id