-
Notifications
You must be signed in to change notification settings - Fork 58
/
Endpoint.java
139 lines (126 loc) · 4 KB
/
Endpoint.java
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
/*
* Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
* with the License. A copy of the License is located at
*
* http://aws.amazon.com/apache2.0/
*
* or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
* OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
* and limitations under the License.
*/
package ai.djl.serving.models;
import ai.djl.serving.workflow.Workflow;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;
/** A class that represents a webservice endpoint. */
public class Endpoint {
private List<Workflow> workflows;
private Map<String, Integer> map;
private AtomicInteger position;
/** Constructs an {@code Endpoint} instance. */
public Endpoint() {
workflows = new ArrayList<>();
map = new ConcurrentHashMap<>();
position = new AtomicInteger(0);
}
/**
* Adds a workflow to the entpoint.
*
* @param workflow the workflow to be added
* @return true if add success
*/
public synchronized boolean add(Workflow workflow) {
String version = workflow.getVersion();
if (version == null) {
if (workflows.isEmpty()) {
map.put("default", 0);
return workflows.add(workflow);
}
return false;
}
if (map.containsKey(version)) {
return false;
}
map.put(version, workflows.size());
return workflows.add(workflow);
}
/**
* Returns the {@link Workflow}s associated with the endpoint.
*
* @return the {@link Workflow}s associated with the endpoint
*/
public List<Workflow> getWorkflows() {
return workflows;
}
/**
* Removes a workflow version from the {@code Endpoint}.
*
* @param version the workflow version
* @return null if the specified version doesn't exist
*/
public synchronized Workflow remove(String version) {
if (version == null) {
if (workflows.isEmpty()) {
return null;
}
Workflow workflow = workflows.remove(0);
reIndex();
return workflow;
}
Integer index = map.remove(version);
if (index == null) {
return null;
}
Workflow workflow = workflows.remove((int) index);
reIndex();
return workflow;
}
/**
* Returns the {@code Workflow} for the specified version.
*
* @param version the version of the workflow to retrieve
* @return the {@code Workflow} for the specified version
*/
public Workflow get(String version) {
Integer index = map.get(version);
if (index == null) {
return null;
}
return workflows.get(index);
}
/**
* Returns the next version of workflow to serve the inference request.
*
* @return the next version of workflow to serve the inference request
*/
public Workflow next() {
int size = workflows.size();
if (size == 1) {
return workflows.get(0);
}
int index = position.getAndUpdate(operand -> (operand + 1) % size);
return workflows.get(index);
}
private void reIndex() {
map.clear();
int size = workflows.size();
for (int i = 0; i < size; ++i) {
Workflow workflow = workflows.get(i);
String version = workflow.getVersion();
if (version != null) {
map.put(version, i);
}
}
}
/** Closes the {@code Endpoint}. */
public void close() {
for (Workflow workflow : workflows) {
workflow.close();
}
}
}