Skip to content

Commit 934a6dc

Browse files
tjaketjake
authored
Ipcexample (#127)
Adds IPCService exmple for basic service over unix sockets
1 parent a3edcbe commit 934a6dc

File tree

8 files changed

+754
-44
lines changed

8 files changed

+754
-44
lines changed

‎jvector-base/src/main/java/io/github/jbellis/jvector/pq/ProductQuantization.java

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import io.github.jbellis.jvector.disk.Io;
2020
import io.github.jbellis.jvector.disk.RandomAccessReader;
2121
import io.github.jbellis.jvector.graph.RandomAccessVectorValues;
22+
import io.github.jbellis.jvector.util.PoolingSupport;
2223
import io.github.jbellis.jvector.util.RamUsageEstimator;
2324
import io.github.jbellis.jvector.util.PhysicalCoreExecutor;
2425
import io.github.jbellis.jvector.vector.VectorUtil;
@@ -60,12 +61,16 @@ public class ProductQuantization {
6061
public static ProductQuantization compute(RandomAccessVectorValues<float[]> ravv, int M, boolean globallyCenter) {
6162
// limit the number of vectors we train on
6263
var P = min(1.0f, MAX_PQ_TRAINING_SET_SIZE / (float) ravv.size());
64+
var ravvCopy = ravv.isValueShared() ? PoolingSupport.newThreadBased(ravv::copy) : PoolingSupport.newNoPooling(ravv);
6365
var subvectorSizesAndOffsets = getSubvectorSizesAndOffsets(ravv.dimension(), M);
6466
var vectors = IntStream.range(0, ravv.size()).parallel()
6567
.filter(i -> ThreadLocalRandom.current().nextFloat() < P)
6668
.mapToObj(targetOrd -> {
67-
float[] v = ravv.vectorValue(targetOrd);
68-
return ravv.isValueShared() ? Arrays.copyOf(v, v.length) : v;
69+
try (var pooledRavv = ravvCopy.get()) {
70+
var localRavv = pooledRavv.get();
71+
float[] v = localRavv.vectorValue(targetOrd);
72+
return localRavv.isValueShared() ? Arrays.copyOf(v, v.length) : v;
73+
}
6974
})
7075
.collect(Collectors.toList());
7176

‎jvector-base/src/main/java/io/github/jbellis/jvector/util/PoolingSupport.java

Lines changed: 74 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package io.github.jbellis.jvector.util;
22

3+
import java.util.Objects;
34
import java.util.concurrent.LinkedBlockingQueue;
45
import java.util.concurrent.atomic.AtomicInteger;
56
import java.util.function.Supplier;
@@ -14,23 +15,14 @@
1415
public abstract class PoolingSupport<T> {
1516

1617
/**
17-
* Creates a pool of objects intended to be used by a fixed thread pool.
18-
* The pool size will the processor count.
18+
* Creates a pool of objects intended to be used by a thread pool.
19+
* This is a replacement for ThreadLocal.
1920
* @param initialValue allows creation of new instances for the pool
2021
*/
2122
public static <T> PoolingSupport<T> newThreadBased(Supplier<T> initialValue) {
2223
return new ThreadPooling<>(initialValue);
2324
}
2425

25-
/**
26-
* Creates a pool intended to be used by a fixed thread pool
27-
* @param threadLimit the specific number of threads to be sharing the pooled objects
28-
* @param initialValue allows creation of new instances for the pool
29-
*/
30-
public static <T> PoolingSupport<T> newThreadBased(int threadLimit, Supplier<T> initialValue) {
31-
return new ThreadPooling<>(threadLimit, initialValue);
32-
}
33-
3426
/**
3527
* Special case of not actually needing a pool (when other times you do)
3628
*
@@ -41,6 +33,16 @@ public static <T> PoolingSupport<T> newNoPooling(T fixedValue) {
4133
}
4234

4335

36+
/**
37+
* Recycling of objects using a MPMC queue
38+
*
39+
* @param limit the specific number of threads to be sharing the pooled objects
40+
* @param initialValue allows creation of new instances for the pool
41+
*/
42+
public static <T> PoolingSupport<T> newQueuePooling(int limit, Supplier<T> initialValue) {
43+
return new QueuedPooling<>(limit, initialValue);
44+
}
45+
4446
private PoolingSupport() {
4547
}
4648

@@ -62,7 +64,7 @@ private PoolingSupport() {
6264
* Internal call used when pooled item is returned
6365
* @param value
6466
*/
65-
protected abstract void onClosed(T value);
67+
protected abstract void onClosed(Pooled<T> value);
6668

6769
/**
6870
* Wrapper class for items in the pool
@@ -71,7 +73,7 @@ private PoolingSupport() {
7173
* in a try-with-resources statement.
7274
* @param <T>
7375
*/
74-
public static class Pooled<T> implements AutoCloseable {
76+
public final static class Pooled<T> implements AutoCloseable {
7577
private final T value;
7678
private final PoolingSupport<T> owner;
7779
private Pooled(PoolingSupport<T> owner, T value) {
@@ -85,56 +87,45 @@ public T get() {
8587

8688
@Override
8789
public void close() {
88-
owner.onClosed(this.value);
90+
owner.onClosed(this);
8991
}
9092
}
9193

9294

93-
static class ThreadPooling<T> extends PoolingSupport<T>
95+
final static class ThreadPooling<T> extends PoolingSupport<T>
9496
{
95-
private final int limit;
96-
private final AtomicInteger created;
97-
private final LinkedBlockingQueue<T> queue;
97+
private final ThreadLocal<Pooled<T>> threadLocal;
9898
private final Supplier<T> initialValue;
9999

100100
private ThreadPooling(Supplier<T> initialValue) {
101-
//+1 for main thread
102-
this(Runtime.getRuntime().availableProcessors() + 1, initialValue);
103-
}
104-
105-
private ThreadPooling(int threadLimit, Supplier<T> initialValue) {
106-
this.limit = threadLimit;
107-
this.created = new AtomicInteger(0);
108-
this.queue = new LinkedBlockingQueue<>(threadLimit);
109101
this.initialValue = initialValue;
102+
this.threadLocal = new ThreadLocal<>();
110103
}
111104

105+
@Override
112106
public Pooled<T> get() {
113-
T t = queue.poll();
114-
if (t != null)
115-
return new Pooled<>(this, t);
107+
Pooled<T> val = threadLocal.get();
108+
if (val != null)
109+
return val;
116110

117-
if (created.incrementAndGet() > limit) {
118-
created.decrementAndGet();
119-
throw new IllegalStateException("Number of outstanding pooled objects has gone beyond the limit of " + limit);
120-
}
121-
return new Pooled<>(this, initialValue.get());
111+
val = new Pooled<>(this, initialValue.get());
112+
threadLocal.set(val);
113+
return val;
122114
}
123115

116+
@Override
124117
public Stream<T> stream() {
125-
if (queue.size() < created.get())
126-
throw new IllegalStateException("close() was not called on all pooled objects yet");
127-
128-
return queue.stream();
118+
throw new UnsupportedOperationException();
129119
}
130120

131-
protected void onClosed(T value) {
132-
queue.offer(value);
121+
@Override
122+
protected void onClosed(Pooled<T> value) {
123+
133124
}
134125
}
135126

136127

137-
static class NoPooling<T> extends PoolingSupport<T> {
128+
final static class NoPooling<T> extends PoolingSupport<T> {
138129
private final T value;
139130
private final Pooled<T> staticPooled;
140131
private NoPooling(T value) {
@@ -153,7 +144,48 @@ public Stream<T> stream() {
153144
}
154145

155146
@Override
156-
protected void onClosed(T value) {
147+
protected void onClosed(Pooled<T> value) {
148+
}
149+
}
150+
151+
152+
final static class QueuedPooling<T> extends PoolingSupport<T> {
153+
private final int limit;
154+
private final AtomicInteger created;
155+
private final LinkedBlockingQueue<Pooled<T>> queue;
156+
private final Supplier<T> initialValue;
157+
158+
private QueuedPooling(int limit, Supplier<T> initialValue) {
159+
this.limit = limit;
160+
this.created = new AtomicInteger(0);
161+
this.queue = new LinkedBlockingQueue<>(limit);
162+
this.initialValue = initialValue;
163+
}
164+
165+
@Override
166+
public Pooled<T> get() {
167+
Pooled<T> t = queue.poll();
168+
if (t != null)
169+
return t;
170+
171+
if (created.incrementAndGet() > limit) {
172+
created.decrementAndGet();
173+
throw new IllegalStateException("Number of outstanding pooled objects has gone beyond the limit of " + limit);
174+
}
175+
return new Pooled<>(this, initialValue.get());
176+
}
177+
178+
@Override
179+
public Stream<T> stream() {
180+
if (queue.size() < created.get())
181+
throw new IllegalStateException("close() was not called on all pooled objects yet");
182+
183+
return queue.stream().filter(Objects::nonNull).map(Pooled::get);
184+
}
185+
186+
@Override
187+
protected void onClosed(Pooled<T> value) {
188+
queue.offer(value);
157189
}
158190
}
159191
}

‎jvector-examples/README.md

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
# JVector Examples
2+
3+
JVector comes with the following sample programs to try:
4+
5+
### SiftSmall
6+
A simple benchmark for the sift dataset located in the [siftsmall](./siftsmall) directory in the project root.
7+
8+
> `mvn compile exec:exec@sift`
9+
10+
### Bench
11+
Performs grid search across the `GraphIndexBuilder` parameter space to find
12+
the best tradeoffs between recall and throughput.
13+
14+
This benchmark requires datasets from [https://github.com/erikbern/ann-benchmarks](https://github.com/erikbern/ann-benchmarks/blob/main/README.md#data-sets) to be downloaded to hdf5 and fvec
15+
directories `hdf5` or `fvec` under the project root depending on the dataset format.
16+
17+
You can use [`plot_output.py`](./plot_output.py) to graph the [pareto-optimal points](https://en.wikipedia.org/wiki/Pareto_efficiency) found by `Bench`.
18+
19+
> `mvn compile exec:exec@bench`
20+
21+
Some sample KNN datasets for testing based on ada-002 embeddings generated on wikipedia data are available in ivec/fvec format for testing at:
22+
23+
```
24+
aws s3 ls s3://astra-vector/wikipedia/ --no-sign-request
25+
PRE 100k/
26+
PRE 1M/
27+
PRE 4M/
28+
```
29+
30+
download them with the aws s3 cli as follows:
31+
32+
```
33+
aws s3 sync s3://astra-vector/wikipedia/100k ./ --no-sign-request
34+
```
35+
36+
To run `SiftSmall`/`Bench` without the JVM vector module available, you can use the following invocations:
37+
38+
> `mvn -Pjdk11 compile exec:exec@bench`
39+
40+
> `mvn -Pjdk11 compile exec:exec@sift`
41+
42+
### IPCService
43+
44+
A simple service for adding / querying vectors over a unix socket.
45+
46+
Install [socat]() using homebrew on mac or apt/rpm on linux
47+
48+
Mac:
49+
> `brew install socat`
50+
51+
Linux:
52+
> `apt-get install socat`
53+
54+
Start the service with:
55+
> `mvn compile exec:exec@ipcserve`
56+
57+
Now you can interact with the service
58+
```bash
59+
socat - unix-client:/tmp/jvector.sock
60+
61+
CREATE 3 DOT_PRODUCT 1 20
62+
OK
63+
WRITE [0.1,0.15,0.3]
64+
OK
65+
WRITE [0.2,0.83,0.05]
66+
OK
67+
WRITE [0.5,0.5,0.5]
68+
OK
69+
OPTIMIZE
70+
OK
71+
SEARCH 20 3 [0.15,0.1,0.1]
72+
RESULT [2,1,0]
73+
```
74+
75+
#### Commands
76+
All commands are completed with `\n`.
77+
78+
No spaces are allowed inside vector brackets.
79+
80+
* `CREATE {dimensions} {similarity-function} {M} {EFConstruction}`
81+
* Creates a new index for this session
82+
83+
* `WRITE [N,N,N] ... [N,N,N]`
84+
* Add one or more vectors to the index
85+
* `OPTIMIZE`
86+
* Call when indexing is complete
87+
* `MEMORY`
88+
* Get the in memory size of index
89+
* `SEARCH {EFSearch} {top-k} [N,N,N] ... [N,N,N]`
90+
* Search index for the top-k closest vectors (ordinals of indexed values returned per query)
91+
* `BULKLOAD {localpath}`
92+
* Bulk loads a local file in numpy format Rows x Columns
93+

‎jvector-examples/pom.xml

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,12 @@
3939
<artifactId>util-mmap</artifactId>
4040
<version>1.0.52-3042601</version>
4141
</dependency>
42+
<dependency>
43+
<groupId>com.kohlschutter.junixsocket</groupId>
44+
<artifactId>junixsocket-core</artifactId>
45+
<version>2.8.1</version>
46+
<type>pom</type>
47+
</dependency>
4248
</dependencies>
4349
<profiles>
4450
<profile>
@@ -127,6 +133,32 @@
127133
</arguments>
128134
</configuration>
129135
</execution>
136+
<execution>
137+
<id>ipcserve</id>
138+
<configuration>
139+
<arguments>
140+
<argument>-classpath</argument>
141+
<classpath/>
142+
<argument>--add-modules=jdk.incubator.vector</argument>
143+
<argument>-ea</argument>
144+
<argument>io.github.jbellis.jvector.example.IPCService</argument>
145+
</arguments>
146+
</configuration>
147+
</execution>
148+
<execution>
149+
<id>ipcserve-1core</id>
150+
<configuration>
151+
<arguments>
152+
<argument>-classpath</argument>
153+
<classpath/>
154+
<argument>--add-modules=jdk.incubator.vector</argument>
155+
<argument>-Djava.util.concurrent.ForkJoinPool.common.parallelism=1</argument>
156+
<argument>-Djvector.physical_core_count=1</argument>
157+
<argument>-ea</argument>
158+
<argument>io.github.jbellis.jvector.example.IPCService</argument>
159+
</arguments>
160+
</configuration>
161+
</execution>
130162
</executions>
131163
</plugin>
132164
</plugins>

0 commit comments

Comments
 (0)