Skip to content

new: [dependencies] split ML into optional extras, reduce default image from 12GB to 593MB#7

Open
wolffhechel wants to merge 1 commit into
vulnerability-lookup:mainfrom
wolffhechel:fix/optional-ml-deps
Open

new: [dependencies] split ML into optional extras, reduce default image from 12GB to 593MB#7
wolffhechel wants to merge 1 commit into
vulnerability-lookup:mainfrom
wolffhechel:fix/optional-ml-deps

Conversation

@wolffhechel

Copy link
Copy Markdown

Summary

Split ML dependencies (torch, transformers) into an optional extras group. The default install provides all 14 API-backed tools without PyTorch. ML classification tools (classify_severity, classify_cwe) are available when installed with -E ml.

Default Docker image: 593 MB (down from 12.1 GB).

Changes

  • pyproject.toml: ML dependencies moved to [tool.poetry.extras]
  • server.py: graceful registration - skips ML tools when torch is unavailable
  • severity.py: absorbs classify_cwe from cwe.py (both ML-dependent)
  • cwe.py: retains only get_recent_vulnerabilities_by_cwe (API-backed, no ML)
  • Dockerfile: Python 3.12, three-mode ML install via INSTALL_ML build arg
  • README.md: documents optional ML install, Docker variants, nvidia-docker
  • .dockerignore: added

Install paths

poetry install              # 14 API tools, no torch (~500 MB)
poetry install -E ml        # 14 API tools + 2 ML classifiers
docker build .              # 593 MB, API tools only
docker build --build-arg INSTALL_ML=cpu .   # ~1.6 GB, CPU-only torch
docker build --build-arg INSTALL_ML=gpu .   # ~5+ GB, CUDA torch

The gpu variant requires nvidia-container-toolkit and docker run --gpus all.

Testing

  • Default install: 14 tools registered, no torch/transformers imported
  • ML install: 16 tools registered, classifiers functional with CPU torch
  • Docker default: 593 MB image, healthy, 14 tools
  • Docker ML (cpu): CPU-only torch, no nvidia-*/triton packages

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant