Reorganize unit tests and update CICD (#37)

* Reorganize unit tests and update CICD

* Fix pylint

Author: JoeZijunZhou

.github/workflows/release.yaml

@@ -31,7 +31,6 @@ jobs:
   release:
     name: Create Release with tag
     runs-on: ${{ matrix.os }}
-    needs: release
     strategy:
       matrix:
         os: [ubuntu-20.04]

.github/workflows/unit_tests.yaml

@@ -74,15 +74,9 @@ jobs:
     - name: Install Dependencies
       run: |
         pip install -r requirements.txt
-    - name: Test mock JetStream token utils
+    - name: Run all unit tests in JetStream (jetstream/tests)
       run: |
-        python -m jetstream.engine.utils_test
-    - name: Test mock JetStream engine implementation
+        coverage run -m unittest -v
+    - name: Create test coverage report
       run: |
-        python -m jetstream.engine.mock_engine_test
-    - name: Test JetStream core orchestrator
-      run: |
-        python -m jetstream.core.orchestrator_test
-    - name: Test JetStream core server library
-      run: |
-        python -m jetstream.core.server_test
+        coverage report -m

.gitignore

@@ -3,6 +3,7 @@ __pycache__
 build/
 dist/
 google_jetstream.egg-info/
+.coverage
 
 # local folders
 data/

jetstream/tests/__init__.py

@@ -0,0 +1,13 @@
+# 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.

jetstream/tests/core/__init__.py

@@ -0,0 +1,13 @@
+# 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.

jetstream/core/orchestrator_test.py renamed to jetstream/tests/core/test_orchestrator.py

@@ -0,0 +1,13 @@
+# 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.

jetstream/core/server_test.py renamed to jetstream/tests/core/test_server.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+"""Tests functionality of the tokenizer with supported models."""
+
 import os
 import unittest
 from typing import List
@@ -24,7 +26,8 @@ class SPTokenizer:
   """Tokenier used in original llama2 git"""
 
   def __init__(self, tokenizer_path: str):
-    self.tokenizer = SentencePieceProcessor(model_file=tokenizer_path)
+    self.tokenizer = SentencePieceProcessor()
+    self.tokenizer.Load(model_file=tokenizer_path)
     assert self.tokenizer.vocab_size() == self.tokenizer.get_piece_size()
 
   def decode(self, t: List[int]) -> str: