[ML] Make inference timeout test more reliable

x-pack/plugin/ml/qa/native-multi-node-tests/src/javaRestTest/java/org/elasticsearch/xpack/ml/integration/PyTorchModelIT.java

@@ -200,8 +200,15 @@ public void testEvaluateWithMinimalTimeout() throws IOException {
         putModelDefinition(modelId);
         putVocabulary(List.of("these", "are", "my", "words"), modelId);
         startDeployment(modelId);
-        ResponseException ex = expectThrows(ResponseException.class, () -> infer("my words", modelId, TimeValue.ZERO));
-        assertThat(ex.getResponse().getStatusLine().getStatusCode(), equalTo(429));
+        // There is a race between inference and timeout so that
+        // even with a zero timeout a valid inference response may
+        // be returned.
+        // The test asserts that if an error occurs it is a timeout error
+        try {
+            infer("my words", modelId, TimeValue.ZERO);
+        } catch (ResponseException ex) {
+            assertThat(ex.getResponse().getStatusLine().getStatusCode(), equalTo(408));
+        }
         stopDeployment(modelId);
     }
 

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/deployment/DeploymentManager.java

@@ -297,7 +297,7 @@ void onTimeout() {
             if (notified.compareAndSet(false, true)) {
                 processContext.getResultProcessor().ignoreResposeWithoutNotifying(String.valueOf(requestId));
                 listener.onFailure(
-                    new ElasticsearchStatusException("timeout [{}] waiting for inference result", RestStatus.TOO_MANY_REQUESTS, timeout)
+                    new ElasticsearchStatusException("timeout [{}] waiting for inference result", RestStatus.REQUEST_TIMEOUT, timeout)
                 );
                 return;
             }