refactor(fluentbit): cleanup protocol lambdas and address codex comments

Author: DragonStuff

src/dstack/_internal/server/services/logs/fluentbit.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Protocol, runtime_checkable
+from typing import List, Optional, Protocol
 from uuid import UUID
 
 from dstack._internal.core.models.logs import (
@@ -21,16 +21,13 @@
 logger = get_logger(__name__)
 
 
-# Check if elasticsearch is available (optional for ship-only mode)
 ELASTICSEARCH_AVAILABLE = True
 try:
     from elasticsearch import Elasticsearch
     from elasticsearch.exceptions import ApiError, TransportError
 except ImportError:
     ELASTICSEARCH_AVAILABLE = False
 else:
-
-    # Catch both API errors and transport/connection errors
     ElasticsearchError: tuple = (ApiError, TransportError)  # type: ignore[misc]
 
     class ElasticsearchReader:
@@ -51,9 +48,7 @@ def __init__(
             try:
                 self._client.info()
             except ElasticsearchError as e:
-                raise LogStorageError(
-                    f"Failed to connect to Elasticsearch/OpenSearch: {e}"
-                ) from e
+                raise LogStorageError(f"Failed to connect to Elasticsearch/OpenSearch: {e}") from e
 
         def read(
             self,
@@ -82,12 +77,16 @@ def read(
             search_params: dict = {
                 "index": self._index,
                 "query": query,
-                "sort": [{"@timestamp": {"order": sort_order}}],
+                "sort": [
+                    {"@timestamp": {"order": sort_order}},
+                    {"_id": {"order": sort_order}},
+                ],
                 "size": request.limit,
             }
 
             if request.next_token:
-                search_params["search_after"] = [request.next_token]
+                parts = request.next_token.split(":", 1)
+                search_params["search_after"] = [parts[0], parts[1]]
 
             try:
                 response = self._client.search(**search_params)
@@ -97,7 +96,7 @@ def read(
 
             hits = response.get("hits", {}).get("hits", [])
             logs = []
-            last_sort_value = None
+            last_sort_values = None
 
             for hit in hits:
                 source = hit.get("_source", {})
@@ -108,9 +107,7 @@ def read(
                     from datetime import datetime
 
                     try:
-                        timestamp = datetime.fromisoformat(
-                            timestamp_str.replace("Z", "+00:00")
-                        )
+                        timestamp = datetime.fromisoformat(timestamp_str.replace("Z", "+00:00"))
                     except ValueError:
                         continue
                 else:
@@ -125,12 +122,12 @@ def read(
                 )
 
                 sort_values = hit.get("sort")
-                if sort_values:
-                    last_sort_value = sort_values[0]
+                if sort_values and len(sort_values) >= 2:
+                    last_sort_values = sort_values
 
             next_token = None
-            if len(logs) == request.limit and last_sort_value is not None:
-                next_token = str(last_sort_value)
+            if len(logs) == request.limit and last_sort_values is not None:
+                next_token = f"{last_sort_values[0]}:{last_sort_values[1]}"
 
             return JobSubmissionLogs(
                 logs=logs,
@@ -149,29 +146,13 @@ def close(self) -> None:
     FLUENTBIT_AVAILABLE = False
 else:
 
-    @runtime_checkable
     class FluentBitWriter(Protocol):
-        """Protocol for Fluent-bit log writers."""
-
-        def write(self, tag: str, records: List[dict]) -> None:
-            """Write log records to Fluent-bit."""
-            ...
+        def write(self, tag: str, records: List[dict]) -> None: ...
+        def close(self) -> None: ...
 
-        def close(self) -> None:
-            """Close any resources."""
-            ...
-
-    @runtime_checkable
     class LogReader(Protocol):
-        """Protocol for log readers (Interface Segregation Principle)."""
-
-        def read(self, stream_name: str, request: PollLogsRequest) -> JobSubmissionLogs:
-            """Read logs from the storage backend."""
-            ...
-
-        def close(self) -> None:
-            """Close any resources."""
-            ...
+        def read(self, stream_name: str, request: PollLogsRequest) -> JobSubmissionLogs: ...
+        def close(self) -> None: ...
 
     class HTTPFluentBitWriter:
         """Writes logs to Fluent-bit via HTTP POST."""
@@ -183,11 +164,21 @@ def __init__(self, host: str, port: int) -> None:
         def write(self, tag: str, records: List[dict]) -> None:
             for record in records:
                 try:
-                    self._client.post(
+                    response = self._client.post(
                         f"{self._endpoint}/{tag}",
                         json=record,
                         headers={"Content-Type": "application/json"},
                     )
+                    response.raise_for_status()
+                except httpx.HTTPStatusError as e:
+                    logger.error(
+                        "Fluent-bit HTTP request failed with status %d: %s",
+                        e.response.status_code,
+                        e.response.text,
+                    )
+                    raise LogStorageError(
+                        f"Fluent-bit HTTP error: status {e.response.status_code}"
+                    ) from e
                 except httpx.HTTPError as e:
                     logger.error("Failed to write log to Fluent-bit via HTTP: %s", e)
                     raise LogStorageError(f"Fluent-bit HTTP error: {e}") from e
@@ -257,7 +248,6 @@ def __init__(
             else:
                 raise LogStorageError(f"Unsupported Fluent-bit protocol: {protocol}")
 
-            # Initialize reader based on configuration (Dependency Inversion Principle)
             self._reader: LogReader
             if es_host:
                 if not ELASTICSEARCH_AVAILABLE:

src/tests/_internal/server/services/test_fluentbit_logs.py

@@ -83,7 +83,33 @@ def test_write_posts_records(self, mock_httpx_client):
             headers={"Content-Type": "application/json"},
         )
 
-    def test_write_raises_on_http_error(self, mock_httpx_client):
+    def test_write_calls_raise_for_status(self, mock_httpx_client):
+        """Test that response.raise_for_status() is called to detect non-2xx responses."""
+        mock_response = Mock()
+        mock_httpx_client.post.return_value = mock_response
+        writer = HTTPFluentBitWriter(host="localhost", port=8080)
+
+        writer.write(tag="test-tag", records=[{"message": "test"}])
+
+        mock_response.raise_for_status.assert_called_once()
+
+    def test_write_raises_on_http_status_error(self, mock_httpx_client):
+        """Test that 4xx/5xx responses are properly detected and raise LogStorageError."""
+        import httpx
+
+        mock_response = Mock()
+        mock_response.status_code = 500
+        mock_response.text = "Internal Server Error"
+        mock_httpx_client.post.return_value = mock_response
+        mock_response.raise_for_status.side_effect = httpx.HTTPStatusError(
+            "Server Error", request=Mock(), response=mock_response
+        )
+        writer = HTTPFluentBitWriter(host="localhost", port=8080)
+
+        with pytest.raises(LogStorageError, match="Fluent-bit HTTP error: status 500"):
+            writer.write(tag="test-tag", records=[{"message": "test"}])
+
+    def test_write_raises_on_transport_error(self, mock_httpx_client):
         import httpx
 
         mock_httpx_client.post.side_effect = httpx.HTTPError("Connection failed")
@@ -468,15 +494,15 @@ def test_read_returns_logs(self, mock_es_client):
                             "message": "Hello",
                             "stream": "test-stream",
                         },
-                        "sort": [1696586513234],
+                        "sort": [1696586513234, "doc1"],
                     },
                     {
                         "_source": {
                             "@timestamp": "2023-10-06T10:01:53.235000+00:00",
                             "message": "World",
                             "stream": "test-stream",
                         },
-                        "sort": [1696586513235],
+                        "sort": [1696586513235, "doc2"],
                     },
                 ]
             }
@@ -499,7 +525,7 @@ def test_read_returns_logs(self, mock_es_client):
             assert len(result.logs) == 2
             assert result.logs[0].message == "Hello"
             assert result.logs[1].message == "World"
-            assert result.next_token == "1696586513235"
+            assert result.next_token == "1696586513235:doc2"
 
     def test_read_with_time_filtering(self, mock_es_client):
         with patch("dstack._internal.server.services.logs.fluentbit.Elasticsearch") as mock:
@@ -540,7 +566,10 @@ def test_read_descending_order(self, mock_es_client):
             reader.read("test-stream", request)
 
             call_args = mock_es_client.search.call_args
-            assert call_args.kwargs["sort"] == [{"@timestamp": {"order": "desc"}}]
+            assert call_args.kwargs["sort"] == [
+                {"@timestamp": {"order": "desc"}},
+                {"_id": {"order": "desc"}},
+            ]
 
     def test_read_with_next_token(self, mock_es_client):
         with patch("dstack._internal.server.services.logs.fluentbit.Elasticsearch") as mock:
@@ -553,13 +582,13 @@ def test_read_with_next_token(self, mock_es_client):
             request = PollLogsRequest(
                 run_name="test-run",
                 job_submission_id=UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e"),
-                next_token="1696586513234",
+                next_token="1696586513234:doc1",
                 limit=100,
             )
             reader.read("test-stream", request)
 
             call_args = mock_es_client.search.call_args
-            assert call_args.kwargs["search_after"] == ["1696586513234"]
+            assert call_args.kwargs["search_after"] == ["1696586513234", "doc1"]
 
     def test_close_closes_client(self, mock_es_client):
         with patch("dstack._internal.server.services.logs.fluentbit.Elasticsearch") as mock: