fix: ruff formatting and allow-direct-references for handler-base dep
This commit is contained in:
@@ -8,6 +8,7 @@ Bridges NATS events to workflow engines using handler-base:
|
|||||||
3. Monitor execution and publish status updates
|
3. Monitor execution and publish status updates
|
||||||
4. Publish completion to "ai.pipeline.status.{request_id}"
|
4. Publish completion to "ai.pipeline.status.{request_id}"
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import logging
|
import logging
|
||||||
from typing import Any, Optional
|
from typing import Any, Optional
|
||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
@@ -137,9 +138,7 @@ class PipelineBridge(Handler):
|
|||||||
|
|
||||||
try:
|
try:
|
||||||
if engine == "argo":
|
if engine == "argo":
|
||||||
run_id = await self._submit_argo(
|
run_id = await self._submit_argo(pipeline["template"], parameters, request_id)
|
||||||
pipeline["template"], parameters, request_id
|
|
||||||
)
|
|
||||||
else:
|
else:
|
||||||
run_id = await self._submit_kubeflow(
|
run_id = await self._submit_kubeflow(
|
||||||
pipeline["pipeline_id"], parameters, request_id
|
pipeline["pipeline_id"], parameters, request_id
|
||||||
@@ -155,9 +154,7 @@ class PipelineBridge(Handler):
|
|||||||
}
|
}
|
||||||
|
|
||||||
# Publish status update
|
# Publish status update
|
||||||
await self.nats.publish(
|
await self.nats.publish(f"ai.pipeline.status.{request_id}", result)
|
||||||
f"ai.pipeline.status.{request_id}", result
|
|
||||||
)
|
|
||||||
|
|
||||||
logger.info(f"Pipeline {pipeline_name} submitted: {run_id}")
|
logger.info(f"Pipeline {pipeline_name} submitted: {run_id}")
|
||||||
return result
|
return result
|
||||||
@@ -170,9 +167,7 @@ class PipelineBridge(Handler):
|
|||||||
"error": str(e),
|
"error": str(e),
|
||||||
}
|
}
|
||||||
|
|
||||||
async def _submit_argo(
|
async def _submit_argo(self, template: str, parameters: dict, request_id: str) -> str:
|
||||||
self, template: str, parameters: dict, request_id: str
|
|
||||||
) -> str:
|
|
||||||
"""Submit workflow to Argo Workflows."""
|
"""Submit workflow to Argo Workflows."""
|
||||||
with create_span("pipeline.submit.argo") as span:
|
with create_span("pipeline.submit.argo") as span:
|
||||||
if span:
|
if span:
|
||||||
@@ -191,10 +186,7 @@ class PipelineBridge(Handler):
|
|||||||
"spec": {
|
"spec": {
|
||||||
"workflowTemplateRef": {"name": template},
|
"workflowTemplateRef": {"name": template},
|
||||||
"arguments": {
|
"arguments": {
|
||||||
"parameters": [
|
"parameters": [{"name": k, "value": str(v)} for k, v in parameters.items()]
|
||||||
{"name": k, "value": str(v)}
|
|
||||||
for k, v in parameters.items()
|
|
||||||
]
|
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
@@ -208,9 +200,7 @@ class PipelineBridge(Handler):
|
|||||||
result = response.json()
|
result = response.json()
|
||||||
return result["metadata"]["name"]
|
return result["metadata"]["name"]
|
||||||
|
|
||||||
async def _submit_kubeflow(
|
async def _submit_kubeflow(self, pipeline_id: str, parameters: dict, request_id: str) -> str:
|
||||||
self, pipeline_id: str, parameters: dict, request_id: str
|
|
||||||
) -> str:
|
|
||||||
"""Submit run to Kubeflow Pipelines."""
|
"""Submit run to Kubeflow Pipelines."""
|
||||||
with create_span("pipeline.submit.kubeflow") as span:
|
with create_span("pipeline.submit.kubeflow") as span:
|
||||||
if span:
|
if span:
|
||||||
@@ -220,10 +210,7 @@ class PipelineBridge(Handler):
|
|||||||
"name": f"{pipeline_id}-{request_id[:8]}",
|
"name": f"{pipeline_id}-{request_id[:8]}",
|
||||||
"pipeline_spec": {
|
"pipeline_spec": {
|
||||||
"pipeline_id": pipeline_id,
|
"pipeline_id": pipeline_id,
|
||||||
"parameters": [
|
"parameters": [{"name": k, "value": str(v)} for k, v in parameters.items()],
|
||||||
{"name": k, "value": str(v)}
|
|
||||||
for k, v in parameters.items()
|
|
||||||
],
|
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -24,6 +24,9 @@ dev = [
|
|||||||
requires = ["hatchling"]
|
requires = ["hatchling"]
|
||||||
build-backend = "hatchling.build"
|
build-backend = "hatchling.build"
|
||||||
|
|
||||||
|
[tool.hatch.metadata]
|
||||||
|
allow-direct-references = true
|
||||||
|
|
||||||
[tool.hatch.build.targets.wheel]
|
[tool.hatch.build.targets.wheel]
|
||||||
packages = ["."]
|
packages = ["."]
|
||||||
only-include = ["pipeline_bridge.py"]
|
only-include = ["pipeline_bridge.py"]
|
||||||
|
|||||||
@@ -1,9 +1,10 @@
|
|||||||
"""
|
"""
|
||||||
Pytest configuration and fixtures for pipeline-bridge tests.
|
Pytest configuration and fixtures for pipeline-bridge tests.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
import os
|
import os
|
||||||
from unittest.mock import AsyncMock, MagicMock, patch
|
from unittest.mock import MagicMock
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
|
|||||||
@@ -1,6 +1,7 @@
|
|||||||
"""
|
"""
|
||||||
Unit tests for PipelineBridge handler.
|
Unit tests for PipelineBridge handler.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
from unittest.mock import AsyncMock, MagicMock, patch
|
from unittest.mock import AsyncMock, MagicMock, patch
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user