Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Inference API] Fix tests in TransportInferenceActionTests #121302

Merged
merged 5 commits into from
Jan 31, 2025
Merged

[Inference API] Fix tests in TransportInferenceActionTests #121302

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
615f856
Fix [CI] TransportInferenceActionTests testRerouting_HandlesTransport…
timgrein Jan 30, 2025
4a564ab
spotlessApply
timgrein Jan 30, 2025
9abc79f
Merge branch 'main' into fix-muted-inference-api-test
timgrein Jan 30, 2025
7ef5fbc
Unmute test
timgrein Jan 30, 2025
dd66ef5
Unmute second test in TransportInferenceActionTests
timgrein Jan 30, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 0 additions & 6 deletions muted-tests.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -362,12 +362,6 @@ tests:
- class: org.elasticsearch.xpack.security.CoreWithSecurityClientYamlTestSuiteIT
method: test {yaml=indices.get_alias/10_basic/Get aliases via /*/_alias/}
issue: https://github.com/elastic/elasticsearch/issues/121290
- class: org.elasticsearch.xpack.inference.action.TransportInferenceActionTests
method: testRerouting_HandlesTransportException_FromOtherNode
issue: https://github.com/elastic/elasticsearch/issues/121292
- class: org.elasticsearch.xpack.inference.action.TransportInferenceActionTests
method: testRerouting_ToOtherNode
issue: https://github.com/elastic/elasticsearch/issues/121293
- class: org.elasticsearch.xpack.inference.common.InferenceServiceNodeLocalRateLimitCalculatorTests
issue: https://github.com/elastic/elasticsearch/issues/121294
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should we unmute this one too? Or is that a separate issue?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a separate issue. I've ran the test with my changes and the reproduction command and it still failed.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

PR for the other CI issue: #121379

- class: org.elasticsearch.entitlement.runtime.policy.PolicyManagerTests
Expand Down
5 changes: 0 additions & 5 deletions .../src/main/java/org/elasticsearch/xpack/inference/action/BaseTransportInferenceAction.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,6 @@

import static org.elasticsearch.core.Strings.format;
import static org.elasticsearch.xpack.inference.InferencePlugin.INFERENCE_API_FEATURE;
import static org.elasticsearch.xpack.inference.common.InferenceAPIClusterAwareRateLimitingFeature.INFERENCE_API_CLUSTER_AWARE_RATE_LIMITING_FEATURE_FLAG;
import static org.elasticsearch.xpack.inference.telemetry.InferenceStats.modelAttributes;
import static org.elasticsearch.xpack.inference.telemetry.InferenceStats.responseAttributes;

Expand Down Expand Up @@ -188,10 +187,6 @@ private void validateRequest(Request request, UnparsedModel unparsedModel) {
}

private NodeRoutingDecision determineRouting(String serviceName, Request request, UnparsedModel unparsedModel) {
if (INFERENCE_API_CLUSTER_AWARE_RATE_LIMITING_FEATURE_FLAG.isEnabled() == false) {
return NodeRoutingDecision.handleLocally();
}

var modelTaskType = unparsedModel.taskType();

// Rerouting not supported or request was already rerouted
Expand Down
10 changes: 5 additions & 5 deletions ...t/java/org/elasticsearch/xpack/inference/action/BaseTransportInferenceActionTestCase.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@
import org.elasticsearch.xpack.core.inference.action.InferenceAction;
import org.elasticsearch.xpack.inference.InferencePlugin;
import org.elasticsearch.xpack.inference.action.task.StreamingTaskManager;
import org.elasticsearch.xpack.inference.common.InferenceServiceNodeLocalRateLimitCalculator;
import org.elasticsearch.xpack.inference.common.InferenceServiceRateLimitCalculator;
import org.elasticsearch.xpack.inference.registry.ModelRegistry;
import org.elasticsearch.xpack.inference.telemetry.InferenceStats;
import org.junit.Before;
Expand Down Expand Up @@ -64,7 +64,7 @@ public abstract class BaseTransportInferenceActionTestCase<Request extends BaseI
protected static final String inferenceId = "inferenceEntityId";
protected InferenceServiceRegistry serviceRegistry;
protected InferenceStats inferenceStats;
protected InferenceServiceNodeLocalRateLimitCalculator inferenceServiceNodeLocalRateLimitCalculator;
protected InferenceServiceRateLimitCalculator inferenceServiceRateLimitCalculator;
protected TransportService transportService;
protected NodeClient nodeClient;

Expand All @@ -79,7 +79,7 @@ public void setUp() throws Exception {
ThreadPool threadPool = mock();
nodeClient = mock();
transportService = mock();
inferenceServiceNodeLocalRateLimitCalculator = mock();
inferenceServiceRateLimitCalculator = mock();
licenseState = mock();
modelRegistry = mock();
serviceRegistry = mock();
Expand All @@ -94,7 +94,7 @@ public void setUp() throws Exception {
serviceRegistry,
inferenceStats,
streamingTaskManager,
inferenceServiceNodeLocalRateLimitCalculator,
inferenceServiceRateLimitCalculator,
nodeClient,
threadPool
);
Expand All @@ -110,7 +110,7 @@ protected abstract BaseTransportInferenceAction<Request> createAction(
InferenceServiceRegistry serviceRegistry,
InferenceStats inferenceStats,
StreamingTaskManager streamingTaskManager,
InferenceServiceNodeLocalRateLimitCalculator inferenceServiceNodeLocalRateLimitCalculator,
InferenceServiceRateLimitCalculator inferenceServiceNodeLocalRateLimitCalculator,
NodeClient nodeClient,
ThreadPool threadPool
);
Expand Down
30 changes: 17 additions & 13 deletions ...src/test/java/org/elasticsearch/xpack/inference/action/TransportInferenceActionTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
import org.elasticsearch.transport.TransportService;
import org.elasticsearch.xpack.core.inference.action.InferenceAction;
import org.elasticsearch.xpack.inference.action.task.StreamingTaskManager;
import org.elasticsearch.xpack.inference.common.InferenceServiceNodeLocalRateLimitCalculator;
import org.elasticsearch.xpack.inference.common.InferenceServiceRateLimitCalculator;
import org.elasticsearch.xpack.inference.common.RateLimitAssignment;
import org.elasticsearch.xpack.inference.registry.ModelRegistry;
import org.elasticsearch.xpack.inference.telemetry.InferenceStats;
Expand Down Expand Up @@ -50,7 +50,7 @@ protected BaseTransportInferenceAction<InferenceAction.Request> createAction(
InferenceServiceRegistry serviceRegistry,
InferenceStats inferenceStats,
StreamingTaskManager streamingTaskManager,
InferenceServiceNodeLocalRateLimitCalculator inferenceServiceNodeLocalRateLimitCalculator,
InferenceServiceRateLimitCalculator inferenceServiceNodeLocalRateLimitCalculator,
NodeClient nodeClient,
ThreadPool threadPool
) {
Expand All @@ -77,7 +77,7 @@ public void testNoRerouting_WhenTaskTypeNotSupported() {
TaskType unsupportedTaskType = TaskType.COMPLETION;
mockService(listener -> listener.onResponse(mock()));

when(inferenceServiceNodeLocalRateLimitCalculator.isTaskTypeReroutingSupported(serviceId, unsupportedTaskType)).thenReturn(false);
when(inferenceServiceRateLimitCalculator.isTaskTypeReroutingSupported(serviceId, unsupportedTaskType)).thenReturn(false);

var listener = doExecute(unsupportedTaskType);

Expand All @@ -89,8 +89,8 @@ public void testNoRerouting_WhenTaskTypeNotSupported() {
public void testNoRerouting_WhenNoGroupingCalculatedYet() {
mockService(listener -> listener.onResponse(mock()));

when(inferenceServiceNodeLocalRateLimitCalculator.isTaskTypeReroutingSupported(serviceId, taskType)).thenReturn(true);
when(inferenceServiceNodeLocalRateLimitCalculator.getRateLimitAssignment(serviceId, taskType)).thenReturn(null);
when(inferenceServiceRateLimitCalculator.isTaskTypeReroutingSupported(serviceId, taskType)).thenReturn(true);
when(inferenceServiceRateLimitCalculator.getRateLimitAssignment(serviceId, taskType)).thenReturn(null);

var listener = doExecute(taskType);

Expand All @@ -102,8 +102,8 @@ public void testNoRerouting_WhenNoGroupingCalculatedYet() {
public void testNoRerouting_WhenEmptyNodeList() {
mockService(listener -> listener.onResponse(mock()));

when(inferenceServiceNodeLocalRateLimitCalculator.isTaskTypeReroutingSupported(serviceId, taskType)).thenReturn(true);
when(inferenceServiceNodeLocalRateLimitCalculator.getRateLimitAssignment(serviceId, taskType)).thenReturn(
when(inferenceServiceRateLimitCalculator.isTaskTypeReroutingSupported(serviceId, taskType)).thenReturn(true);
when(inferenceServiceRateLimitCalculator.getRateLimitAssignment(serviceId, taskType)).thenReturn(
new RateLimitAssignment(List.of())
);

Expand All @@ -120,10 +120,10 @@ public void testRerouting_ToOtherNode() {

// The local node is different to the "other-node" responsible for serviceId
when(nodeClient.getLocalNodeId()).thenReturn("local-node");
when(inferenceServiceNodeLocalRateLimitCalculator.isTaskTypeReroutingSupported(serviceId, taskType)).thenReturn(true);
when(inferenceServiceRateLimitCalculator.isTaskTypeReroutingSupported(serviceId, taskType)).thenReturn(true);
// Requests for serviceId are always routed to "other-node"
var assignment = new RateLimitAssignment(List.of(otherNode));
when(inferenceServiceNodeLocalRateLimitCalculator.getRateLimitAssignment(serviceId, taskType)).thenReturn(assignment);
when(inferenceServiceRateLimitCalculator.getRateLimitAssignment(serviceId, taskType)).thenReturn(assignment);

mockService(listener -> listener.onResponse(mock()));
var listener = doExecute(taskType);
Expand All @@ -141,9 +141,9 @@ public void testRerouting_ToLocalNode_WithoutGoingThroughTransportLayerAgain() {

// The local node is the only one responsible for serviceId
when(nodeClient.getLocalNodeId()).thenReturn(localNodeId);
when(inferenceServiceNodeLocalRateLimitCalculator.isTaskTypeReroutingSupported(serviceId, taskType)).thenReturn(true);
when(inferenceServiceRateLimitCalculator.isTaskTypeReroutingSupported(serviceId, taskType)).thenReturn(true);
var assignment = new RateLimitAssignment(List.of(localNode));
when(inferenceServiceNodeLocalRateLimitCalculator.getRateLimitAssignment(serviceId, taskType)).thenReturn(assignment);
when(inferenceServiceRateLimitCalculator.getRateLimitAssignment(serviceId, taskType)).thenReturn(assignment);

mockService(listener -> listener.onResponse(mock()));
var listener = doExecute(taskType);
Expand All @@ -158,9 +158,9 @@ public void testRerouting_HandlesTransportException_FromOtherNode() {
when(otherNode.getId()).thenReturn("other-node");

when(nodeClient.getLocalNodeId()).thenReturn("local-node");
when(inferenceServiceNodeLocalRateLimitCalculator.isTaskTypeReroutingSupported(serviceId, taskType)).thenReturn(true);
when(inferenceServiceRateLimitCalculator.isTaskTypeReroutingSupported(serviceId, taskType)).thenReturn(true);
var assignment = new RateLimitAssignment(List.of(otherNode));
when(inferenceServiceNodeLocalRateLimitCalculator.getRateLimitAssignment(serviceId, taskType)).thenReturn(assignment);
when(inferenceServiceRateLimitCalculator.getRateLimitAssignment(serviceId, taskType)).thenReturn(assignment);

mockService(listener -> listener.onResponse(mock()));

Expand All @@ -173,6 +173,10 @@ public void testRerouting_HandlesTransportException_FromOtherNode() {

var listener = doExecute(taskType);

// Verify request was rerouted
verify(transportService).sendRequest(same(otherNode), eq(InferenceAction.NAME), any(), any());
// Verify local execution didn't happen
verify(listener, never()).onResponse(any());
// Verify exception was propagated from "other-node" to "local-node"
verify(listener).onFailure(same(expectedException));
}
Expand Down
6 changes: 3 additions & 3 deletions .../java/org/elasticsearch/xpack/inference/action/TransportUnifiedCompletionActionTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
import org.elasticsearch.transport.TransportService;
import org.elasticsearch.xpack.core.inference.action.UnifiedCompletionAction;
import org.elasticsearch.xpack.inference.action.task.StreamingTaskManager;
import org.elasticsearch.xpack.inference.common.InferenceServiceNodeLocalRateLimitCalculator;
import org.elasticsearch.xpack.inference.common.InferenceServiceRateLimitCalculator;
import org.elasticsearch.xpack.inference.registry.ModelRegistry;
import org.elasticsearch.xpack.inference.telemetry.InferenceStats;

Expand Down Expand Up @@ -49,7 +49,7 @@ protected BaseTransportInferenceAction<UnifiedCompletionAction.Request> createAc
InferenceServiceRegistry serviceRegistry,
InferenceStats inferenceStats,
StreamingTaskManager streamingTaskManager,
InferenceServiceNodeLocalRateLimitCalculator inferenceServiceNodeLocalRateLimitCalculator,
InferenceServiceRateLimitCalculator inferenceServiceRateLimitCalculator,
NodeClient nodeClient,
ThreadPool threadPool
) {
Expand All @@ -61,7 +61,7 @@ protected BaseTransportInferenceAction<UnifiedCompletionAction.Request> createAc
serviceRegistry,
inferenceStats,
streamingTaskManager,
inferenceServiceNodeLocalRateLimitCalculator,
inferenceServiceRateLimitCalculator,
nodeClient,
threadPool
);
Expand Down