Integrate category classifier model in tracker (#2756)

Signed-off-by: Sergio Castaño Arteaga <[email protected]>

Author: tegioz

.gitattributes

@@ -3,4 +3,4 @@ web/build/** linguist-generated
 __fixtures__/** linguist-generated
 *.sql linguist-detectable=true
 *.sql linguist-language=sql
-ml/categories/model/** linguist-generated
+ml/category/model/** linguist-generated

.github/workflows/ci.yml

@@ -13,6 +13,12 @@ jobs:
         uses: actions/setup-go@v3
         with:
           go-version: 1.19
+      - name: Install TensorFlow C library
+        run: |
+          FILENAME=libtensorflow-cpu-linux-x86_64-2.11.0.tar.gz
+          wget -q --no-check-certificate https://storage.googleapis.com/tensorflow/libtensorflow/${FILENAME}
+          sudo tar -C /usr/local -xzf ${FILENAME}
+          sudo ldconfig /usr/local/lib
       - name: Run golangci-lint
         uses: golangci/golangci-lint-action@v3
         with:
@@ -98,6 +104,12 @@ jobs:
           key: ${{ runner.os }}-go-${{ hashFiles('**/go.sum') }}
           restore-keys: |
             ${{ runner.os }}-go-
+      - name: Install TensorFlow C library
+        run: |
+          FILENAME=libtensorflow-cpu-linux-x86_64-2.11.0.tar.gz
+          wget -q --no-check-certificate https://storage.googleapis.com/tensorflow/libtensorflow/${FILENAME}
+          sudo tar -C /usr/local -xzf ${FILENAME}
+          sudo ldconfig /usr/local/lib
       - name: Run backend tests
         run: go test -cover -race -v -mod=readonly ./...
 
@@ -171,6 +183,12 @@ jobs:
           key: ${{ runner.os }}-go-${{ hashFiles('**/go.sum') }}
           restore-keys: |
             ${{ runner.os }}-go-
+      - name: Install TensorFlow C library
+        run: |
+          FILENAME=libtensorflow-cpu-linux-x86_64-2.11.0.tar.gz
+          wget -q --no-check-certificate https://storage.googleapis.com/tensorflow/libtensorflow/${FILENAME}
+          sudo tar -C /usr/local -xzf ${FILENAME}
+          sudo ldconfig /usr/local/lib
       - name: Build hub
         working-directory: ./cmd/hub
         run: go build -v

.gitignore

@@ -19,4 +19,4 @@ docs/www/content/topics/*
 docs/www/content/topics/annotations/*
 !docs/www/content/topics/annotations/_index.md
 dist
-ml/categories/data/generated
+ml/category/data/generated

charts/artifact-hub/Chart.yaml

@@ -2,7 +2,7 @@ apiVersion: v2
 name: artifact-hub
 description: Artifact Hub is a web-based application that enables finding, installing, and publishing Kubernetes packages.
 type: application
-version: 1.12.1-3
+version: 1.12.1-4
 appVersion: 1.12.0
 kubeVersion: ">= 1.19.0-0"
 home: https://artifacthub.io

charts/artifact-hub/templates/tracker_secret.yaml

@@ -28,3 +28,4 @@ stringData:
       repositoriesNames: {{ .Values.tracker.repositoriesNames }}
       repositoriesKinds: {{ .Values.tracker.repositoriesKinds }}
       bypassDigestCheck: {{ .Values.tracker.bypassDigestCheck }}
+      categoryModelPath: ./ml/category/model

cmd/scanner/Dockerfile

@@ -10,7 +10,7 @@ RUN CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -o /scanner .
 # Trivy installer
 FROM alpine:3.17.1 AS trivy-installer
 RUN apk --no-cache add curl
-RUN curl -sfL https://raw.githubusercontent.com/aquasecurity/trivy/master/contrib/install.sh | sh -s -- -b /usr/local/bin v0.36.1
+RUN curl -sfL https://raw.githubusercontent.com/aquasecurity/trivy/main/contrib/install.sh | sh -s -- -b /usr/local/bin v0.36.1
 
 # Final stage
 FROM alpine:3.17.1

cmd/tracker/Dockerfile

@@ -1,22 +1,39 @@
 # Build tracker
-FROM golang:1.19.5-alpine3.17 AS builder
+FROM golang:1.19.5-bullseye AS builder
+WORKDIR /tmp
+ENV LIBTENSORFLOW_TGZ libtensorflow-cpu-linux-x86_64-2.11.0.tar.gz
+RUN wget -q --no-check-certificate https://storage.googleapis.com/tensorflow/libtensorflow/$LIBTENSORFLOW_TGZ
+RUN tar -C /usr/local -xzf $LIBTENSORFLOW_TGZ
+RUN ldconfig /usr/local/lib
 WORKDIR /go/src/github.com/artifacthub/hub
 COPY go.* ./
 COPY cmd/tracker cmd/tracker
 COPY internal internal
 WORKDIR /go/src/github.com/artifacthub/hub/cmd/tracker
-RUN CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -o /tracker .
+RUN CGO_ENABLED=1 GOOS=linux GOARCH=amd64 go build -o /tracker .
 
 # OPM installer
-FROM golang:1.17-alpine3.16 AS opm-installer
-RUN apk --no-cache add build-base
+FROM golang:1.17-bullseye AS opm-installer
 RUN GO111MODULE=on go get github.com/operator-framework/operator-registry/cmd/[email protected]
 
 # Final stage
-FROM alpine:3.17.1
-RUN apk --no-cache add ca-certificates && addgroup -S tracker -g 1000 && adduser -S tracker -u 1000 -G tracker
+FROM debian:bullseye-slim
+RUN apt-get update \
+    && apt-get install -y ca-certificates \
+    && groupadd -g 1000 tracker \
+    && useradd -u 1000 -g tracker tracker
+WORKDIR /tmp
+ENV LIBTENSORFLOW_TGZ libtensorflow-cpu-linux-x86_64-2.11.0.tar.gz
+RUN apt-get install -y wget \
+    && export LIBTENSORFLOW_TGZ=libtensorflow-cpu-linux-x86_64-2.11.0.tar.gz \
+    && wget -q --no-check-certificate https://storage.googleapis.com/tensorflow/libtensorflow/$LIBTENSORFLOW_TGZ \
+    && tar -C /usr/local -xzf $LIBTENSORFLOW_TGZ \
+    && rm $LIBTENSORFLOW_TGZ \
+    && apt-get remove -y wget \
+    && ldconfig /usr/local/lib
 USER 1000
 WORKDIR /home/tracker
+COPY ml ./ml
 COPY --from=builder /tracker ./
 COPY --from=opm-installer /go/bin/opm /usr/local/bin
 CMD ["./tracker"]

cmd/tracker/main.go

@@ -70,6 +70,7 @@ func main() {
 	}
 	ec := repo.NewErrorsCollector(rm, repo.Tracker)
 	op := oci.NewPuller(cfg)
+	pcc := tracker.NewPackageCategoryClassifierML(cfg.GetString("tracker.categoryModelPath"))
 	svc := &hub.TrackerServices{
 		Ctx:                ctx,
 		Cfg:                cfg,
@@ -82,6 +83,7 @@ func main() {
 		Op:                 op,
 		Is:                 is,
 		Sc:                 oci.NewSignatureChecker(cfg, op),
+		Pcc:                pcc,
 		SetupTrackerSource: tracker.SetupSource,
 	}
 

database/migrations/functions/packages/register_package.sql

@@ -66,6 +66,7 @@ begin
         is_operator,
         channels,
         default_channel,
+        package_category_id,
         repository_id
     ) values (
         v_name,
@@ -75,6 +76,7 @@ begin
         (p_pkg->>'is_operator')::boolean,
         nullif(p_pkg->'channels', 'null'),
         nullif(p_pkg->>'default_channel', ''),
+        nullif((p_pkg->>'category')::int, 0),
         v_repository_id
     )
     on conflict (repository_id, name) do update
@@ -85,7 +87,8 @@ begin
         tsdoc = generate_package_tsdoc(v_name, v_alternative_name, v_display_name, v_description, v_keywords, v_ts_repository, v_ts_publisher),
         is_operator = excluded.is_operator,
         channels = excluded.channels,
-        default_channel = excluded.default_channel
+        default_channel = excluded.default_channel,
+        package_category_id = excluded.package_category_id
     where is_latest(
         v_repository_kind_id,
         v_version,

database/migrations/schema/045_package_category.sql

@@ -0,0 +1,21 @@
+create table if not exists package_category (
+    package_category_id integer primary key,
+    name text not null check (name <> ''),
+    display_name text not null check (display_name <> '')
+);
+
+insert into package_category values (1, 'ai-machine-learning', 'AI / Machine learning');
+insert into package_category values (2, 'database', 'Database');
+insert into package_category values (3, 'integration-delivery', 'Integration and delivery');
+insert into package_category values (4, 'monitoring-logging', 'Monitoring and logging');
+insert into package_category values (5, 'networking', 'Networking');
+insert into package_category values (6, 'security', 'Security');
+insert into package_category values (7, 'storage', 'Storage');
+insert into package_category values (8, 'streaming-messaging', 'Streaming and messaging');
+
+alter table package add column package_category_id integer;
+
+---- create above / drop below ----
+
+alter table package drop column package_category_id;
+drop table if exists package_category;