Add original job definition name to pipeline metrics

Flared · Oct 17, 2024 · dc96a9e · dc96a9e
1 parent cdf69ea
commit dc96a9e
Show file tree

Hide file tree

Showing 8 changed files with 30 additions and 2 deletions.
diff --git a/src/saturn_engine/core/api.py b/src/saturn_engine/core/api.py
@@ -73,6 +73,7 @@ class QueueItem:
     config: dict[str, Any] = field(default_factory=dict)
     labels: dict[str, str] = field(default_factory=dict)
     executor: str = "default"
+    job_definition_name: str = ""
 
     def with_state(self, state: "QueueItemState") -> "QueueItemWithState":
         return QueueItemWithState(**self.__dict__, state=state)

diff --git a/src/saturn_engine/worker/services/metrics.py b/src/saturn_engine/worker/services/metrics.py
@@ -14,6 +14,7 @@
 
 def executable_params(xmsg: ExecutableMessage) -> dict:
     return {
+        "saturn.job_definition.name": xmsg.queue.definition.job_definition_name,
         "saturn.executor.name": xmsg.queue.definition.executor,
         "saturn.job.name": xmsg.queue.definition.name,
         "pipeline": xmsg.message.info.name,

diff --git a/src/saturn_engine/worker/services/usage_metrics.py b/src/saturn_engine/worker/services/usage_metrics.py
@@ -19,6 +19,7 @@
 
 class PipelineName(t.NamedTuple):
     executor: str
+    job_definition_name: str
     name: str
 
 
@@ -38,6 +39,7 @@ class PipelineName(t.NamedTuple):
 
 class PipelineUsage(t.NamedTuple):
     executor: str
+    job_definition_name: str
     name: str
     usage: float
 
@@ -88,6 +90,7 @@ class StageState(t.Generic[T, U]):
     def _name(xmsg: ExecutableMessage) -> PipelineName:
         return PipelineName(
             executor=xmsg.queue.definition.executor or "default",
+            job_definition_name=xmsg.queue.definition.job_definition_name or "",
             name=xmsg.message.info.name,
         )
 
@@ -104,7 +107,10 @@ def pop(self, xmsg: ExecutableMessage) -> U:
     def collect(self, *, now: Nanoseconds) -> t.Iterator[PipelineUsage]:
         for k, state in self.pipelines.items():
             yield PipelineUsage(
-                executor=k.executor, name=k.name, usage=state.collect(now=now)
+                job_definition_name=k.job_definition_name,
+                executor=k.executor,
+                name=k.name,
+                usage=state.collect(now=now),
             )
 
 
@@ -175,6 +181,7 @@ def collect_usage_metrics(
                 yield Observation(
                     pipeline.usage,
                     {
+                        "job_definition_name": pipeline.job_definition_name,
                         "executor": pipeline.executor,
                         "pipeline": pipeline.name,
                         "state": stage_name,

diff --git a/src/saturn_engine/worker_manager/config/declarative_job.py b/src/saturn_engine/worker_manager/config/declarative_job.py
@@ -99,6 +99,7 @@ def to_core_objects(
                 labels=labels,
                 config=self.config,
                 executor=self.executor,
+                job_definition_name=name,
             )
 
 

diff --git a/tests/worker/services/test_metrics.py b/tests/worker/services/test_metrics.py
@@ -17,6 +17,7 @@ def fake_published_message() -> Mock:
     data = Mock()
     data.xmsg.queue.definition.executor = "exec"
     data.xmsg.queue.definition.name = "test-job"
+    data.xmsg.queue.definition.job_definition_name = "test-job-definition"
     data.xmsg.queue.definition.labels = {"k": "v"}
     data.xmsg.message.info.name = "test.fake.pipeline"
     data.topic.name = "test.fake.topic"
@@ -30,10 +31,12 @@ async def test_message_metrics(
     data = Mock()
     data.queue.definition.name = "test-job"
     data.queue.definition.executor = "exec"
+    data.queue.definition.job_definition_name = "test-job-definition"
     data.queue.definition.labels = {"k": "v"}
     data.message.info.name = "test.fake.pipeline"
     pipeline_params = {
         "pipeline": data.message.info.name,
+        "saturn.job_definition.name": "test-job-definition",
         "saturn.executor.name": "exec",
         "saturn.job.name": "test-job",
         "saturn.job.labels.k": "v",
@@ -69,10 +72,12 @@ async def test_metrics_message_executed(
     data = Mock()
     data.queue.definition.name = "test-job"
     data.queue.definition.executor = "exec"
+    data.queue.definition.job_definition_name = "test-job-definition"
     data.queue.definition.labels = {"k": "v"}
     data.message.info.name = "test.fake.pipeline"
     pipeline_params = {
         "pipeline": data.message.info.name,
+        "saturn.job_definition.name": "test-job-definition",
         "saturn.executor.name": "exec",
         "saturn.job.name": "test-job",
         "saturn.job.labels.k": "v",
@@ -145,10 +150,12 @@ async def test_metrics_message_execute_failed(
     data = Mock()
     data.queue.definition.name = "test-job"
     data.queue.definition.executor = "exec"
+    data.queue.definition.job_definition_name = "test-job-definition"
     data.queue.definition.labels = {"k": "v"}
     data.message.info.name = "test.fake.pipeline"
     pipeline_params = {
         "pipeline": data.message.info.name,
+        "saturn.job_definition.name": "test-job-definition",
         "saturn.executor.name": "exec",
         "saturn.job.name": "test-job",
         "saturn.job.labels.k": "v",
@@ -198,6 +205,7 @@ async def test_metrics_message_published(
     params = {
         "pipeline": "test.fake.pipeline",
         "topic": "test.fake.topic",
+        "saturn.job_definition.name": "test-job-definition",
         "saturn.executor.name": "exec",
         "saturn.job.name": "test-job",
         "saturn.job.labels.k": "v",
@@ -233,6 +241,7 @@ async def test_metrics_message_publish_failed(
     params = {
         "pipeline": "test.fake.pipeline",
         "topic": "test.fake.topic",
+        "saturn.job_definition.name": "test-job-definition",
         "saturn.executor.name": "exec",
         "saturn.job.name": "test-job",
         "saturn.job.labels.k": "v",
@@ -268,6 +277,7 @@ async def test_metrics_topic_blocked(
     params = {
         "pipeline": "test.fake.pipeline",
         "topic": "test.fake.topic",
+        "saturn.job_definition.name": "test-job-definition",
         "saturn.executor.name": "exec",
         "saturn.job.name": "test-job",
         "saturn.job.labels.k": "v",

diff --git a/tests/worker/services/test_usage_metrics.py b/tests/worker/services/test_usage_metrics.py
@@ -26,10 +26,15 @@ async def test_message_metrics(
         xmsg: ExecutableMessage = Mock()
         xmsg.message.info.name = pipeline_name
         xmsg.queue.definition.executor = "default"
+        xmsg.queue.definition.job_definition_name = "my-job-definition"
         xmsgs.append(xmsg)
 
     any_: t.Any = None
-    pipeline_params = {"pipeline": pipeline_name, "executor": "default"}
+    pipeline_params = {
+        "pipeline": pipeline_name,
+        "executor": "default",
+        "job_definition_name": "my-job-definition",
+    }
     results = PipelineResults(outputs=[], resources=[])
     metric = services_manager.services.cast_service(UsageMetrics)
 

diff --git a/tests/worker/test_broker.py b/tests/worker/test_broker.py
@@ -99,6 +99,7 @@ async def test_broker_dummy(
                 labels={"owner": "team-saturn"},
                 output={},
                 executor="e1",
+                job_definition_name="job-definition-name",
             )
         ],
         resources=[
@@ -159,6 +160,7 @@ async def test_broker_dummy(
 
     # Test metrics
     pipeline_params = {
+        "saturn.job_definition.name": "job-definition-name",
         "saturn.executor.name": "e1",
         "saturn.job.name": "j1",
         "pipeline": "tests.worker.test_broker.pipeline",

diff --git a/tests/worker_manager/api/test_job_definitions.py b/tests/worker_manager/api/test_job_definitions.py
@@ -75,6 +75,7 @@ def test_api_job_definitions_loaded_from_str(
                     },
                     "name": "test-job-definition",
                     "executor": "default",
+                    "job_definition_name": "test-job-definition",
                     "output": {
                         "default": [
                             {