fix

.pre-commit-config.yaml

@@ -1,20 +1,16 @@
 repos:
-- repo: https://github.com/pre-commit/pre-commit-hooks
-  rev: v3.4.0  # Use the latest revision
-  hooks:
-    - id: trailing-whitespace
-    - id: end-of-file-fixer
-    - id: check-yaml
-- repo: https://github.com/psf/black
-  rev: 24.2.0
-  hooks:
-  - id: black
-    name: Format code
-    args:
-      - --line-length=120
-- repo: https://github.com/pycqa/flake8
-  rev: 7.0.0
-  hooks:
-  - id: flake8
-    name: flake8
-    args: ['--max-line-length=120']
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v3.4.0 # Use the latest revision
+    hooks:
+      - id: trailing-whitespace
+      - id: end-of-file-fixer
+      - id: check-yaml
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    # Ruff version.
+    rev: v0.8.3
+    hooks:
+      # Run the linter.
+      - id: ruff
+        args: [--fix]
+      # Run the formatter.
+      - id: ruff-format

server/lorax_server/utils/graph.py

@@ -41,7 +41,7 @@
 CACHED_BATCH_SIZES = [1, 2, 3, 4, 8, 16] + [
     BATCH_SIZE_INCREMENT * (i + 1) for i in range(MAX_BATCH_SIZE // BATCH_SIZE_INCREMENT)
 ]
-CACHED_BATCH_SIZES = [b for b in CACHED_BATCH_SIZES if b <= MAX_BATCH_SIZE]
+CACHED_BATCH_SIZES = [b for b in CACHED_BATCH_SIZES if b <= COMPILE_BATCH_SIZE]
 
 # Include 0 to ensure we can use cuda graphs without adapters
 # TODO(travis): use padding to allow for more ranks without increasing memory usage