Skip to content

Commit f336e98

Browse files
alexjstalexjst
authored
feat(embedding): ship ONNX/DJL only in dedicated ONNX image variants [PFP-3657] (#19753)
1 parent 6ece48b commit f336e98

40 files changed

Lines changed: 1650 additions & 1429 deletions

File tree

‎datahub-graphql-core/gradle.lockfile‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
# This is a Gradle generated file for dependency locking.
22
# Manual edits can break the build and are not advised.
33
# This file is expected to be part of source control.
4-
ai.djl.huggingface:tokenizers:0.36.0=runtimeClasspath,testRuntimeClasspath
5-
ai.djl:api:0.36.0=runtimeClasspath,testRuntimeClasspath
4+
ai.djl.huggingface:tokenizers:0.36.0=compileClasspath,testCompileClasspath
5+
ai.djl:api:0.36.0=compileClasspath,testCompileClasspath
66
antlr:antlr:2.7.7=compileClasspath,runtimeClasspath,testCompileClasspath,testRuntimeClasspath
77
aopalliance:aopalliance:1.0=compileClasspath,runtimeClasspath,testCompileClasspath,testRuntimeClasspath
88
at.yawk.lz4:lz4-java:1.11.1=runtimeClasspath,testRuntimeClasspath
@@ -107,7 +107,7 @@ com.linkedin.pegasus:restli-client:29.74.2=compileClasspath,runtimeClasspath,tes
107107
com.linkedin.pegasus:restli-common:29.74.2=compileClasspath,runtimeClasspath,testCompileClasspath,testRuntimeClasspath
108108
com.linkedin.pegasus:restli-disruptor:29.74.2=compileClasspath,runtimeClasspath,testCompileClasspath,testRuntimeClasspath
109109
com.linkedin.pegasus:restli-server:29.74.2=runtimeClasspath,testRuntimeClasspath
110-
com.microsoft.onnxruntime:onnxruntime:1.29.0=runtimeClasspath,testRuntimeClasspath
110+
com.microsoft.onnxruntime:onnxruntime:1.29.0=compileClasspath,testCompileClasspath
111111
com.mixpanel:mixpanel-java:1.4.4=runtimeClasspath,testRuntimeClasspath
112112
com.sun.activation:jakarta.activation:2.0.1=runtimeClasspath,testRuntimeClasspath
113113
com.sun.activation:javax.activation:1.2.0=compileClasspath,runtimeClasspath,testCompileClasspath,testRuntimeClasspath
@@ -312,8 +312,8 @@ javax.validation:validation-api:2.0.1.Final=compileClasspath,runtimeClasspath,te
312312
joda-time:joda-time:2.12.7=compileClasspath,runtimeClasspath,testCompileClasspath,testRuntimeClasspath
313313
net.bytebuddy:byte-buddy-agent:1.18.11=testCompileClasspath,testRuntimeClasspath
314314
net.bytebuddy:byte-buddy:1.18.11=testCompileClasspath,testRuntimeClasspath
315-
net.java.dev.jna:jna:5.13.0=compileClasspath,testCompileClasspath
316-
net.java.dev.jna:jna:5.17.0=runtimeClasspath,testRuntimeClasspath
315+
net.java.dev.jna:jna:5.13.0=runtimeClasspath,testRuntimeClasspath
316+
net.java.dev.jna:jna:5.17.0=compileClasspath,testCompileClasspath
317317
net.openhft:zero-allocation-hashing:0.7=compileClasspath,runtimeClasspath,testCompileClasspath,testRuntimeClasspath
318318
net.sf.jopt-simple:jopt-simple:5.0.4=compileClasspath,runtimeClasspath,testCompileClasspath,testRuntimeClasspath
319319
nl.basjes.collections:prefixmap:2.0=compileClasspath,runtimeClasspath,testCompileClasspath,testRuntimeClasspath

‎datahub-upgrade/build.gradle‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -143,6 +143,32 @@ bootJar {
143143
archiveFileName = "${project.name}.jar"
144144
}
145145

146+
// ONNX image variant (R3): a second boot archive that re-adds the ONNX/DJL native libs
147+
// (compileOnlyApi in metadata-io, so absent from the default jar) into BOOT-INF/lib. The
148+
// datahub-upgrade-onnx image is built from this; system-update needs it to build the semantic
149+
// ES index via the embedding provider bean. The default datahub-upgrade.jar stays slim.
150+
configurations {
151+
onnxRuntimeVariant {
152+
canBeConsumed = false
153+
canBeResolved = true
154+
extendsFrom configurations.implementation, configurations.runtimeOnly
155+
}
156+
}
157+
dependencies {
158+
onnxRuntimeVariant externalDependency.onnxRuntime
159+
onnxRuntimeVariant externalDependency.djlApi
160+
onnxRuntimeVariant externalDependency.djlHuggingFaceTokenizers
161+
}
162+
tasks.register('onnxBootJar', org.springframework.boot.gradle.tasks.bundling.BootJar) {
163+
mainClass = 'com.linkedin.datahub.upgrade.UpgradeCliApplication'
164+
archiveFileName = 'datahub-upgrade-onnx.jar'
165+
classpath = sourceSets.main.output + configurations.onnxRuntimeVariant
166+
duplicatesStrategy = DuplicatesStrategy.EXCLUDE
167+
// The Spring Boot plugin only auto-populates this on the primary bootJar; a second BootJar
168+
// must set it or the task fails with "property 'targetJavaVersion' has no value available".
169+
targetJavaVersion = java.targetCompatibility
170+
}
171+
146172
bootRun {
147173
environment "ENTITY_REGISTRY_CONFIG_PATH", "../metadata-models/src/main/resources/entity-registry.yml"
148174
environment "ENABLE_STRUCTURED_PROPERTIES_SYSTEM_UPDATE", "true"

‎datahub-upgrade/gradle.lockfile‎

Lines changed: 632 additions & 632 deletions
Large diffs are not rendered by default.

‎docker/build.gradle‎

Lines changed: 43 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -291,6 +291,8 @@ ext {
291291
':metadata-jobs:mae-consumer-job',
292292
':datahub-actions',
293293
':docker:datahub-ingestion',
294+
':docker:datahub-gms-onnx',
295+
':docker:datahub-upgrade-onnx',
294296
]
295297
]
296298
]
@@ -325,6 +327,16 @@ ext {
325327
return activeProfile
326328
}
327329

330+
// Helper function to read the captured quickstart task name from file.
331+
// Distinct from the profile: some debug stacks share a compose profile (quickstartDebug and
332+
// quickstartDebugBuiltinEmbedding both use 'debug'), so reload[Env] needs the exact task to
333+
// rebuild stack-specific artifacts like the onnx boot archives.
334+
readCapturedTaskName = {
335+
def composeFileName = new File(compose_base).getName()
336+
def taskFile = new File(rootProject.buildDir, composeFileName.replace('.yml', '-task.txt'))
337+
return taskFile.exists() ? taskFile.text.trim() : null
338+
}
339+
328340
// Helper function to find the taskName and config that uses a given profile (only searches debug tasks)
329341
// Returns a map with 'taskName' and 'config' keys
330342
findTaskNameByProfile = { profile ->
@@ -661,8 +673,13 @@ quickstart_configs.each { taskName, config ->
661673
}
662674
}
663675

664-
// quickstartDebugBuiltinEmbedding auto-downloads the ONNX model if not present
665-
tasks.getByName("quickstartDebugBuiltinEmbeddingComposeUp").dependsOn(tasks.getByName("prepareOnnxModel"))
676+
// quickstartDebugBuiltinEmbedding auto-downloads the ONNX model if not present, and builds the
677+
// onnx boot archives that docker-compose.onnx-override.yml mounts over the default war/jar so the
678+
// debug stack exercises the real ONNX-variant artifacts.
679+
tasks.getByName("quickstartDebugBuiltinEmbeddingComposeUp").dependsOn(
680+
tasks.getByName("prepareOnnxModel"),
681+
':metadata-service:war:onnxBootJar',
682+
':datahub-upgrade:onnxBootJar')
666683

667684

668685
tasks.register('minDockerCompose2.20', Exec) {
@@ -781,6 +798,12 @@ tasks.withType(ComposeUp).configureEach {
781798
profileFile.text = primaryProfile
782799

783800
logger.lifecycle("Captured profile '${primaryProfile}' to ${profileFile.absolutePath}")
801+
802+
// Also record the exact quickstart task. Profiles are not unique across debug stacks
803+
// (quickstartDebug and quickstartDebugBuiltinEmbedding both use 'debug'), so reload[Env]
804+
// reads this to rebuild stack-specific artifacts like the onnx boot archives.
805+
def taskFile = new File(rootProject.buildDir, composeFileName.replace('.yml', '-task.txt'))
806+
taskFile.text = taskBaseName
784807
}
785808
}
786809
}
@@ -807,6 +830,15 @@ tasks.register("reload", Exec) {
807830

808831
// Dynamically depend on the correct prepareAll task
809832
dependsOn tasks.named("prepareAll${matchingTaskName}")
833+
834+
// The builtin-embedding stack mounts the onnx boot archives (war-onnx.war /
835+
// datahub-upgrade-onnx.jar) over the defaults, so reload must rebuild THOSE or the
836+
// containers keep running stale variant code after a source change. Key off the
837+
// captured task, not matchingTaskName: builtin-embedding shares the 'debug' profile
838+
// with quickstartDebug, so profile lookup resolves to quickstartDebug here.
839+
if (readCapturedTaskName.call() == 'quickstartDebugBuiltinEmbedding') {
840+
dependsOn ':metadata-service:war:onnxBootJar', ':datahub-upgrade:onnxBootJar'
841+
}
810842
}
811843
}
812844

@@ -864,6 +896,15 @@ tasks.register("reloadEnv", Exec) {
864896
matchingConfig = matchingTask.config
865897
// Dynamically depend on the correct prepareAll task
866898
dependsOn tasks.named("prepareAll${matchingTaskName}")
899+
900+
// The builtin-embedding stack mounts the onnx boot archives (war-onnx.war /
901+
// datahub-upgrade-onnx.jar) over the defaults, so reloadEnv must rebuild THOSE or the
902+
// recreated containers keep running stale variant code after a source change. Key off
903+
// the captured task, not matchingTaskName: builtin-embedding shares the 'debug'
904+
// profile with quickstartDebug, so profile lookup resolves to quickstartDebug here.
905+
if (readCapturedTaskName.call() == 'quickstartDebugBuiltinEmbedding') {
906+
dependsOn ':metadata-service:war:onnxBootJar', ':datahub-upgrade:onnxBootJar'
907+
}
867908
}
868909
}
869910

‎docker/datahub-gms-onnx/build.gradle‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
plugins {
2+
id 'java' // required for versioning
3+
}
4+
5+
apply from: "../../gradle/versioning/versioning.gradle"
6+
apply from: "../../gradle/docker/docker.gradle"
7+
8+
// ONNX image variant of datahub-gms (R3): identical to datahub-gms but built from the onnx boot
9+
// archive (war-onnx.war), which bundles the ONNX/DJL native libs. Reuses the datahub-gms Dockerfile
10+
// via the WAR_FILE build arg; the default datahub-gms image stays slim.
11+
ext {
12+
docker_registry = rootProject.ext.docker_registry == 'linkedin' ? 'acryldata' : docker_registry
13+
docker_repo = 'datahub-gms-onnx'
14+
// Reuse the datahub-gms Dockerfile + context (start.sh, prometheus config, etc.).
15+
base_docker_dir = 'datahub-gms'
16+
}
17+
18+
// Ensure the war project is configured first so its onnxBootJar task is registered before we
19+
// reference it below.
20+
evaluationDependsOn(':metadata-service:war')
21+
def onnxWar = project(':metadata-service:war').tasks.named('onnxBootJar')
22+
23+
docker {
24+
dependsOn ':metadata-service:war:onnxBootJar'
25+
name "${docker_registry}/${docker_repo}:${versionTag}"
26+
dockerfile file("${rootProject.projectDir}/docker/${base_docker_dir}/Dockerfile")
27+
files onnxWar.flatMap { it.archiveFile }
28+
files fileTree(rootProject.projectDir) {
29+
include '.dockerignore'
30+
include 'docker/monitoring/*'
31+
include 'docker/snippets/*'
32+
include "docker/${base_docker_dir}/*"
33+
include 'metadata-models/src/main/resources/*'
34+
}.exclude {
35+
i -> (!i.file.name.endsWith(".dockerignore") && i.file.isHidden())
36+
}
37+
38+
def dockerBuildArgs = [WAR_FILE: 'war-onnx.war']
39+
if (project.hasProperty('dockerBaseImage')) {
40+
dockerBuildArgs.BASE_IMAGE = project.getProperty('dockerBaseImage')
41+
}
42+
dockerBuildArgs.APK_REPOSITORY_URL = gradle.ext['datahub.dependencies.apkrepo.url']
43+
dockerBuildArgs.GITHUB_REPO_URL = gradle.ext['datahub.dependencies.github.baseURL']
44+
dockerBuildArgs.MAVEN_CENTRAL_REPO_URL = gradle.ext['datahub.dependencies.maven.central']
45+
buildArgs(dockerBuildArgs)
46+
}

‎docker/datahub-gms-onnx/gradle.lockfile‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
# This is a Gradle generated file for dependency locking.
2+
# Manual edits can break the build and are not advised.
3+
# This file is expected to be part of source control.
4+
jakarta.annotation:jakarta.annotation-api:3.0.0=compileClasspath,runtimeClasspath,testCompileClasspath,testRuntimeClasspath
5+
empty=annotationProcessor,testAnnotationProcessor

‎docker/datahub-gms/Dockerfile‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,11 @@ COPY metadata-models/src/main/resources/entity-registry.yml /datahub/datahub-gms
2828
COPY docker/datahub-gms/start.sh /datahub/datahub-gms/scripts/start.sh
2929
COPY docker/monitoring/client-prometheus-config.yaml /datahub/datahub-gms/scripts/prometheus-config.yaml
3030
RUN chmod +x /datahub/datahub-gms/scripts/start.sh
31-
COPY metadata-service/war/build/libs/war.war /datahub/datahub-gms/bin/war.war
31+
# WAR_FILE selects the boot archive: war.war (default, slim) or war-onnx.war (datahub-gms-onnx
32+
# variant, which bundles the ONNX/DJL native libs). Destination filename stays war.war so start.sh
33+
# is unchanged.
34+
ARG WAR_FILE=war.war
35+
COPY metadata-service/war/build/libs/${WAR_FILE} /datahub/datahub-gms/bin/war.war
3236

3337
FROM base AS dev-install
3438
# Dummy stage for development. Assumes code is built on your machine and mounted to this image.

‎docker/datahub-upgrade-onnx/build.gradle‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
plugins {
2+
id 'java' // required for versioning
3+
}
4+
5+
apply from: "../../gradle/versioning/versioning.gradle"
6+
apply from: "../../gradle/docker/docker.gradle"
7+
8+
// ONNX image variant of datahub-upgrade (R3): identical to datahub-upgrade but built from the onnx
9+
// boot archive (datahub-upgrade-onnx.jar). system-update needs the ONNX/DJL libs to build the
10+
// semantic ES index via the embedding provider bean. Reuses the datahub-upgrade Dockerfile via the
11+
// UPGRADE_JAR build arg; the default datahub-upgrade image stays slim.
12+
ext {
13+
docker_registry = rootProject.ext.docker_registry == 'linkedin' ? 'acryldata' : docker_registry
14+
docker_repo = 'datahub-upgrade-onnx'
15+
base_docker_dir = 'datahub-upgrade'
16+
}
17+
18+
// Ensure the datahub-upgrade project is configured first so its onnxBootJar task is registered
19+
// before we reference it below.
20+
evaluationDependsOn(':datahub-upgrade')
21+
def onnxJar = project(':datahub-upgrade').tasks.named('onnxBootJar')
22+
23+
docker {
24+
dependsOn ':datahub-upgrade:onnxBootJar'
25+
name "${docker_registry}/${docker_repo}:${versionTag}"
26+
dockerfile file("${rootProject.projectDir}/docker/${base_docker_dir}/Dockerfile")
27+
files onnxJar.flatMap { it.archiveFile }
28+
files fileTree(rootProject.projectDir) {
29+
include '.dockerignore'
30+
include 'docker/monitoring/*'
31+
include 'docker/snippets/*'
32+
include "docker/${base_docker_dir}/*"
33+
include 'metadata-models/src/main/resources/*'
34+
}.exclude {
35+
i -> (!i.file.name.endsWith(".dockerignore") && i.file.isHidden())
36+
}
37+
38+
def dockerBuildArgs = [UPGRADE_JAR: 'datahub-upgrade-onnx.jar']
39+
if (project.hasProperty('dockerBaseImage')) {
40+
dockerBuildArgs.BASE_IMAGE = project.getProperty('dockerBaseImage')
41+
}
42+
dockerBuildArgs.APK_REPOSITORY_URL = gradle.ext['datahub.dependencies.apkrepo.url']
43+
dockerBuildArgs.GITHUB_REPO_URL = gradle.ext['datahub.dependencies.github.baseURL']
44+
dockerBuildArgs.MAVEN_CENTRAL_REPO_URL = gradle.ext['datahub.dependencies.maven.central']
45+
buildArgs(dockerBuildArgs)
46+
}

‎docker/datahub-upgrade-onnx/gradle.lockfile‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
# This is a Gradle generated file for dependency locking.
2+
# Manual edits can break the build and are not advised.
3+
# This file is expected to be part of source control.
4+
jakarta.annotation:jakarta.annotation-api:3.0.0=compileClasspath,runtimeClasspath,testCompileClasspath,testRuntimeClasspath
5+
empty=annotationProcessor,testAnnotationProcessor

‎docker/datahub-upgrade/Dockerfile‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,11 @@ RUN --mount=type=cache,target=/var/cache/apk,sharing=locked \
2525
COPY docker/snippets/wait_for_deps.sh /usr/local/lib/datahub/wait_for_deps.sh
2626

2727
FROM base AS prod-install
28-
COPY datahub-upgrade/build/libs/datahub-upgrade.jar /datahub/datahub-upgrade/bin/
28+
# UPGRADE_JAR selects the boot archive: datahub-upgrade.jar (default, slim) or
29+
# datahub-upgrade-onnx.jar (datahub-upgrade-onnx variant, bundling ONNX/DJL). Destination filename
30+
# stays datahub-upgrade.jar so start.sh is unchanged.
31+
ARG UPGRADE_JAR=datahub-upgrade.jar
32+
COPY datahub-upgrade/build/libs/${UPGRADE_JAR} /datahub/datahub-upgrade/bin/datahub-upgrade.jar
2933
COPY docker/datahub-upgrade/start.sh /datahub/datahub-upgrade/scripts/start.sh
3034
COPY docker/monitoring/client-prometheus-config.yaml /datahub/datahub-upgrade/scripts/prometheus-config.yaml
3135
RUN chmod +x /datahub/datahub-upgrade/scripts/start.sh

0 commit comments

Comments
 (0)