Skip to content

Commit 5770fb3

Browse files
Distributed Tree Storage (#8423)
* WIP: distributed skeleton tree storage * add count_skeleton_keys tool * implement loading and saving with external skeleton bodies * changelog, migration guide * migrate trees previously stored in the main tracing * add missing shutdown * remove unused imports in count skeletons script * pr feedback, take care of revertToVersion with distributed skeletons * typo * Fix duplicated line in nmlgenerator * do not try to flush deleted trees * extract the info if skeleton is stored with external tree bodies to single variable. reset it during duplicate and merge * remove hasExternalTreeBody tree property, use centralized one for all * cleanup, unify save method signature * refresh e2e snapshots (new bool in skeletons) * add storedWithExternalTreeBodies to frontend skeleton type * fix js type * use proto index 1 and 2 * add wrapper class SkeletonTracingWithUpdatedTreeIds to avoid the tuple --------- Co-authored-by: MichaelBuessemeyer <[email protected]>
1 parent b427ff2 commit 5770fb3

File tree

22 files changed

+229
-53
lines changed

22 files changed

+229
-53
lines changed

CHANGELOG.unreleased.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ For upgrade instructions, please check the [migration guide](MIGRATIONS.released
1515
- Added a command palette that allows navigating between pages, switching tools and accessing some user settings via Ctrl+P. [#8447](https://github.com/scalableminds/webknossos/pull/8447/)
1616
- Super users can now share the trained AI models with other organizations. [#8418](https://github.com/scalableminds/webknossos/pull/8418)
1717
- Failed jobs may be retried by super-users. [#8377](https://github.com/scalableminds/webknossos/pull/8377)
18+
- Optimized server-side storage of skeleton annotation layers. [#8423](https://github.com/scalableminds/webknossos/pull/8423)
1819

1920
### Changed
2021
- When using a zarr link to a wk-served data layer as another layer’s source, the user’s token is used to access the data. [#8322](https://github.com/scalableminds/webknossos/pull/8322/)

MIGRATIONS.unreleased.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ User-facing changes are documented in the [changelog](CHANGELOG.released.md).
88
## Unreleased
99
[Commits](https://github.com/scalableminds/webknossos/compare/25.02.1...HEAD)
1010

11+
- FossilDB now needs to be opened with additional column family `skeletonTreeBodies`. [#8423](https://github.com/scalableminds/webknossos/pull/8423)
1112

1213
### Postgres Evolutions:
1314
- [126-mag-real-paths.sql](conf/evolutions/126-mag-real-paths.sql)

docker-compose.yml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -277,7 +277,7 @@ services:
277277
command:
278278
- fossildb
279279
- -c
280-
- skeletons,volumes,volumeData,volumeSegmentIndex,editableMappingsInfo,editableMappingsAgglomerateToGraph,editableMappingsSegmentToAgglomerate,annotations,annotationUpdates
280+
- skeletons,skeletonTreeBodies,volumes,volumeData,volumeSegmentIndex,editableMappingsInfo,editableMappingsAgglomerateToGraph,editableMappingsSegmentToAgglomerate,annotations,annotationUpdates
281281
user: ${USER_UID:-fossildb}:${USER_GID:-fossildb}
282282

283283
fossildb-persisted:

fossildb/run.sh

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,6 @@ if [ ! -f "$JAR" ] || [ ! "$CURRENT_VERSION" == "$VERSION" ]; then
1414
wget -q --show-progress -O "$JAR" "$URL"
1515
fi
1616

17-
COLLECTIONS="skeletons,volumes,volumeData,volumeSegmentIndex,editableMappingsInfo,editableMappingsAgglomerateToGraph,editableMappingsSegmentToAgglomerate,annotations,annotationUpdates"
17+
COLLECTIONS="skeletons,skeletonTreeBodies,volumes,volumeData,volumeSegmentIndex,editableMappingsInfo,editableMappingsAgglomerateToGraph,editableMappingsSegmentToAgglomerate,annotations,annotationUpdates"
1818

1919
exec java -jar "$JAR" -c "$COLLECTIONS" -d "$FOSSILDB_HOME/data" -b "$FOSSILDB_HOME/backup"

frontend/javascripts/test/snapshots/public-test/test-bundle/test/backend-snapshot-tests/annotations.e2e.js.md

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1238,6 +1238,7 @@ Generated by [AVA](https://avajs.dev).
12381238
z: 0,
12391239
},
12401240
id: 'id',
1241+
storedWithExternalTreeBodies: true,
12411242
treeGroups: [],
12421243
trees: [],
12431244
typ: 'Skeleton',
@@ -1332,6 +1333,7 @@ Generated by [AVA](https://avajs.dev).
13321333
z: 0,
13331334
},
13341335
id: 'id',
1336+
storedWithExternalTreeBodies: true,
13351337
treeGroups: [],
13361338
trees: [],
13371339
typ: 'Skeleton',
@@ -1421,6 +1423,7 @@ Generated by [AVA](https://avajs.dev).
14211423
z: 3,
14221424
},
14231425
id: 'id',
1426+
storedWithExternalTreeBodies: true,
14241427
treeGroups: [],
14251428
trees: [],
14261429
typ: 'Skeleton',
@@ -1447,6 +1450,7 @@ Generated by [AVA](https://avajs.dev).
14471450
z: 0,
14481451
},
14491452
id: 'id',
1453+
storedWithExternalTreeBodies: true,
14501454
treeGroups: [
14511455
{
14521456
children: [
@@ -2296,6 +2300,7 @@ Generated by [AVA](https://avajs.dev).
22962300
z: 0,
22972301
},
22982302
id: 'id',
2303+
storedWithExternalTreeBodies: true,
22992304
treeGroups: [],
23002305
trees: [
23012306
{

frontend/javascripts/types/api_flow_types.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -879,6 +879,7 @@ export type ServerSkeletonTracing = ServerTracingBase & {
879879
boundingBox?: ServerBoundingBox;
880880
trees: Array<ServerSkeletonTracingTree>;
881881
treeGroups: Array<TreeGroup> | null | undefined;
882+
storedWithExternalTreeBodies?: boolean; // unused in frontend
882883
};
883884
export type ServerVolumeTracing = ServerTracingBase & {
884885
// The following property is added when fetching the
Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
import logging
2+
from utils import setup_logging, log_since
3+
import argparse
4+
from connections import connect_to_fossildb, assert_grpc_success
5+
import time
6+
import fossildbapi_pb2 as proto
7+
8+
logger = logging.getLogger(__name__)
9+
10+
11+
def main():
12+
print("Hello from count_skeleton_keys")
13+
setup_logging()
14+
parser = argparse.ArgumentParser()
15+
parser.add_argument("--src", type=str, help="Source fossildb host and port. Example: localhost:7155", required=True)
16+
args = parser.parse_args()
17+
before = time.time()
18+
src_stub = connect_to_fossildb(args.src, "source")
19+
collection = "skeletons"
20+
list_keys_page_size = 100
21+
key_count = 0
22+
current_start_after_key = None
23+
while True:
24+
print(key_count, end=" ", flush=True)
25+
list_keys_reply = src_stub.ListKeys(proto.ListKeysRequest(collection=collection, limit=list_keys_page_size, startAfterKey=current_start_after_key))
26+
assert_grpc_success(list_keys_reply)
27+
this_page_key_count = len(list_keys_reply.keys)
28+
key_count += this_page_key_count
29+
if this_page_key_count == 0:
30+
# We iterated towards the very end of the collection
31+
break
32+
current_start_after_key = list_keys_reply.keys[-1]
33+
log_since(before, f"Saw {key_count} keys.")
34+
35+
36+
if __name__ == '__main__':
37+
main()

tools/nml/nmlgenerator.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,6 @@ def print_trees(self):
4141
for node_id in range(node_id_count, node_id_count + self.nodes_per_tree):
4242
print(' <node ')
4343
print(f' id="{node_id}" radius="165.0" x="{node_id}" y="{node_id}" z="{tree_id}"')
44-
print(f' id="{node_id}" radius="165.0" x="{node_id}" y="{node_id}" z="{tree_id}"')
4544
print(' inVp="0" inMag="0" bitDepth="8" interpolation="false" time="1395338380800">')
4645
print(' </node>')
4746
print(' </nodes>')

webknossos-datastore/proto/SkeletonTracing.proto

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -38,10 +38,15 @@ enum TreeTypeProto {
3838
AGGLOMERATE = 1;
3939
}
4040

41+
message TreeBody {
42+
repeated Node nodes = 1;
43+
repeated Edge edges = 2;
44+
}
45+
4146
message Tree {
4247
required int32 treeId = 1;
43-
repeated Node nodes = 2;
44-
repeated Edge edges = 3;
48+
repeated Node nodes = 2; // empty for skeletons with storedWithExternalTreeBodies=true
49+
repeated Edge edges = 3; // empty for skeletons with storedWithExternalTreeBodies=true
4550
optional ColorProto color = 4;
4651
repeated BranchPoint branchPoints = 5;
4752
repeated Comment comments = 6;
@@ -77,6 +82,7 @@ message SkeletonTracing {
7782
optional string organizationId = 13; // used when parsing and handling nmls, not used in tracing store anymore, do not rely on correct values
7883
repeated AdditionalCoordinateProto editPositionAdditionalCoordinates = 21;
7984
repeated AdditionalAxisProto additionalAxes = 22; // Additional axes for which this tracing is defined
85+
optional bool storedWithExternalTreeBodies = 23; // Only set to true on storing. If false, all trees will be stored
8086
}
8187

8288
message SkeletonTracingOpt {

webknossos-tracingstore/app/com/scalableminds/webknossos/tracingstore/TSRemoteWebknossosClient.scala

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ import com.scalableminds.webknossos.datastore.services.{
1818
UserAccessRequest
1919
}
2020
import com.scalableminds.webknossos.tracingstore.annotation.AnnotationLayerParameters
21+
import com.scalableminds.webknossos.tracingstore.tracings.skeleton.SkeletonTracingWithUpdatedTreeIds
2122
import com.typesafe.scalalogging.LazyLogging
2223
import play.api.inject.ApplicationLifecycle
2324
import play.api.libs.json.{JsObject, Json, OFormat}
@@ -104,7 +105,7 @@ class TSRemoteWebknossosClient @Inject()(
104105

105106
def createTracingFor(annotationId: String,
106107
layerParameters: AnnotationLayerParameters,
107-
previousVersion: Long): Fox[Either[SkeletonTracing, VolumeTracing]] = {
108+
previousVersion: Long): Fox[Either[SkeletonTracingWithUpdatedTreeIds, VolumeTracing]] = {
108109
val req = rpc(s"$webknossosUri/api/tracingstores/$tracingStoreName/createTracing")
109110
.addQueryString("annotationId" -> annotationId)
110111
.addQueryString("previousVersion" -> previousVersion.toString) // used for fetching old precedence layers
@@ -115,9 +116,12 @@ class TSRemoteWebknossosClient @Inject()(
115116
.postJsonWithProtoResponse[AnnotationLayerParameters, VolumeTracing](layerParameters)(VolumeTracing)
116117
.map(Right(_))
117118
case AnnotationLayerType.Skeleton =>
118-
req
119-
.postJsonWithProtoResponse[AnnotationLayerParameters, SkeletonTracing](layerParameters)(SkeletonTracing)
120-
.map(Left(_))
119+
for {
120+
skeletonTracing <- req.postJsonWithProtoResponse[AnnotationLayerParameters, SkeletonTracing](layerParameters)(
121+
SkeletonTracing)
122+
} yield
123+
Left[SkeletonTracingWithUpdatedTreeIds, VolumeTracing](
124+
SkeletonTracingWithUpdatedTreeIds(skeletonTracing, Set.empty))
121125
}
122126
}
123127

webknossos-tracingstore/app/com/scalableminds/webknossos/tracingstore/annotation/AnnotationWithTracings.scala

Lines changed: 35 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import com.scalableminds.webknossos.tracingstore.tracings.editablemapping.{
1111
EditableMappingUpdateAction,
1212
EditableMappingUpdater
1313
}
14+
import com.scalableminds.webknossos.tracingstore.tracings.skeleton.SkeletonTracingWithUpdatedTreeIds
1415
import com.scalableminds.webknossos.tracingstore.tracings.skeleton.updating.SkeletonUpdateAction
1516
import com.scalableminds.webknossos.tracingstore.tracings.volume.ApplyableVolumeUpdateAction
1617
import com.typesafe.scalalogging.LazyLogging
@@ -20,7 +21,7 @@ import scala.concurrent.ExecutionContext
2021

2122
case class AnnotationWithTracings(
2223
annotation: AnnotationProto,
23-
tracingsById: Map[String, Either[SkeletonTracing, VolumeTracing]],
24+
tracingsById: Map[String, Either[SkeletonTracingWithUpdatedTreeIds, VolumeTracing]],
2425
editableMappingsByTracingId: Map[String, (EditableMappingInfo, EditableMappingUpdater)])
2526
extends LazyLogging {
2627

@@ -32,17 +33,26 @@ case class AnnotationWithTracings(
3233
for {
3334
tracingEither <- tracingsById.get(tracingId)
3435
skeletonTracing <- tracingEither match {
35-
case Left(st: SkeletonTracing) => Full(st)
36-
case _ => Failure(f"Tried to access tracing $tracingId as skeleton, but is volume")
36+
case Left(SkeletonTracingWithUpdatedTreeIds(skeletonTracing, _)) => Full(skeletonTracing)
37+
case _ => Failure(f"Tried to access tracing $tracingId as skeleton, but is volume")
3738
}
3839
} yield skeletonTracing
3940

4041
def getSkeletons: List[(String, SkeletonTracing)] =
4142
tracingsById.view.flatMap {
42-
case (id, Left(st: SkeletonTracing)) => Some(id, st)
43-
case _ => None
43+
case (id, Left(SkeletonTracingWithUpdatedTreeIds(skeletonTracing, _))) => Some(id, skeletonTracing)
44+
case _ => None
4445
}.toList
4546

47+
def getUpdatedTreeBodyIdsForSkeleton(tracingId: String): Box[Set[Int]] =
48+
for {
49+
tracingEither <- tracingsById.get(tracingId)
50+
updatedTreeIds <- tracingEither match {
51+
case Left(SkeletonTracingWithUpdatedTreeIds(_, updatedTreeIds)) => Full(updatedTreeIds)
52+
case _ => Failure(f"Tried to access tracing $tracingId as skeleton to access updated tree ids, but is volume")
53+
}
54+
} yield updatedTreeIds
55+
4656
def getVolumes: List[(String, VolumeTracing)] =
4757
tracingsById.view.flatMap {
4858
case (id, Right(vt: VolumeTracing)) => Some(id, vt)
@@ -86,7 +96,7 @@ case class AnnotationWithTracings(
8696

8797
def addLayer(a: AddLayerAnnotationAction,
8898
tracingId: String,
89-
tracing: Either[SkeletonTracing, VolumeTracing]): AnnotationWithTracings =
99+
tracing: Either[SkeletonTracingWithUpdatedTreeIds, VolumeTracing]): AnnotationWithTracings =
90100
this.copy(
91101
annotation = annotation.copy(
92102
annotationLayers = annotation.annotationLayers :+ AnnotationLayerProto(
@@ -115,8 +125,9 @@ case class AnnotationWithTracings(
115125

116126
def withVersion(newVersion: Long): AnnotationWithTracings = {
117127
val tracingsUpdated = tracingsById.view.mapValues {
118-
case Left(t: SkeletonTracing) => Left(t.withVersion(newVersion))
119-
case Right(t: VolumeTracing) => Right(t.withVersion(newVersion))
128+
case Left(t: SkeletonTracingWithUpdatedTreeIds) =>
129+
Left(t.withVersion(newVersion))
130+
case Right(t: VolumeTracing) => Right(t.withVersion(newVersion))
120131
}
121132
this.copy(
122133
annotation = annotation.copy(version = newVersion,
@@ -142,17 +153,19 @@ case class AnnotationWithTracings(
142153
def applySkeletonAction(a: SkeletonUpdateAction)(implicit ec: ExecutionContext): Fox[AnnotationWithTracings] =
143154
for {
144155
skeletonTracing <- getSkeleton(a.actionTracingId)
156+
previousUpdatedTreeIds <- getUpdatedTreeBodyIdsForSkeleton(a.actionTracingId)
145157
updated = a.applyOn(skeletonTracing)
146-
} yield this.copy(tracingsById = tracingsById.updated(a.actionTracingId, Left(updated)))
158+
newUpdatedTreeIds = previousUpdatedTreeIds.concat(a.updatedTreeBodyIds)
159+
} yield
160+
this.copy(
161+
tracingsById =
162+
tracingsById.updated(a.actionTracingId, Left(SkeletonTracingWithUpdatedTreeIds(updated, newUpdatedTreeIds))))
147163

148164
def applyVolumeAction(a: ApplyableVolumeUpdateAction)(implicit ec: ExecutionContext): Fox[AnnotationWithTracings] =
149165
for {
150166
volumeTracing <- getVolume(a.actionTracingId)
151167
updated = a.applyOn(volumeTracing)
152-
} yield
153-
AnnotationWithTracings(annotation,
154-
tracingsById.updated(a.actionTracingId, Right(updated)),
155-
editableMappingsByTracingId)
168+
} yield this.copy(tracingsById = tracingsById.updated(a.actionTracingId, Right(updated)))
156169

157170
def applyEditableMappingAction(a: EditableMappingUpdateAction)(
158171
implicit ec: ExecutionContext): Fox[AnnotationWithTracings] =
@@ -171,4 +184,13 @@ case class AnnotationWithTracings(
171184
} yield ()
172185
}
173186

187+
def markAllTreeBodiesAsChanged: AnnotationWithTracings = {
188+
val newTracingsById = tracingsById.view.map {
189+
case (id, Left(st: SkeletonTracingWithUpdatedTreeIds)) =>
190+
(id, Left[SkeletonTracingWithUpdatedTreeIds, VolumeTracing](st.markAllTreeBodiesAsChanged))
191+
case other => other
192+
}.toMap
193+
this.copy(tracingsById = newTracingsById)
194+
}
195+
174196
}

0 commit comments

Comments
 (0)