Add unit tests CI github action

.github/workflows/UnitTests.yaml

@@ -0,0 +1,72 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# This workflow will install Python dependencies, run tests and lint with a variety of Python versions
+# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python
+
+name: Unit Tests
+
+on:
+  pull_request:
+  push:
+    branches: [ "main" ]
+  workflow_dispatch:
+  schedule:
+    # Run the job every 12 hours
+    - cron:  '0 */12 * * *'
+
+jobs:
+  py:
+    name: "Python type check"
+    strategy:
+      matrix:
+        os: [ubuntu-20.04]
+        python-version: ['3.10']
+    runs-on: ${{ matrix.os }}
+    steps:
+    - uses: actions/checkout@v4
+    - name: Setup Python
+      uses: actions/setup-python@v4
+      with:
+        python-version: ${{ matrix.python-version }}
+    - name: Install Dependencies
+      run: |
+        pip install pytype
+        pip install -r requirements.txt
+    - name: Typecheck the code with pytype
+      run: |
+        pytype --jobs auto --disable import-error --disable module-attr jetstream/
+
+  cpu:
+    name: "JetStream unit tests"
+    strategy:
+      matrix:
+        os: [ubuntu-20.04]
+        python-version: ['3.10']
+    runs-on: ${{ matrix.os }}
+    steps:
+    - uses: actions/checkout@v4
+    - name: Setup Python
+      uses: actions/setup-python@v4
+      with:
+        python-version: ${{ matrix.python-version }}
+    - name: Install Dependencies
+      run: |
+        pip install -r requirements.txt
+    - name: Test mock JetStream token utils
+      run: |
+        python -m jetstream.engine.utils_test
+    - name: Test mock JetStream engine implementation
+      run: |
+        python -m jetstream.engine.mock_engine_test

jetstream/core/orchestrator_test.py

@@ -61,6 +61,8 @@ def _setup_driver(self):
     driver = orchestrator.Driver(
         prefill_engines=[prefill_engine],
         generate_engines=[generate_engine],
+        prefill_params=[prefill_engine.load_params()],
+        generate_params=[generate_engine.load_params()],
     )
     return driver
 

jetstream/core/server_test.py

@@ -18,7 +18,7 @@
 response.
 """
 
-from typing import Any
+from typing import Any, Type
 
 import grpc
 import portpicker
@@ -48,7 +48,7 @@ class ServerTest(parameterized.TestCase):
   )
   def test_server(
       self,
-      config: config_lib.ServerConfig,
+      config: Type[config_lib.ServerConfig],
       expected_tokens: list[str],
       devices: list[Any],
   ):