Reorganize unit tests and update CICD

.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/tests/engine/__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/engine/test_token_utils.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:

requirements.in

@@ -1,4 +1,5 @@
 absl-py
+coverage
 flax
 grpcio
 jax

requirements.txt

@@ -37,6 +37,8 @@ clu==0.0.10
     # via seqio
 contextlib2==21.6.0
     # via ml-collections
+coverage==7.4.4
+    # via -r requirements.in
 dm-tree==0.1.8
     # via
     #   chex
@@ -49,7 +51,6 @@ etils[array-types,enp,epath,epy,etqdm,etree]==1.6.0
     # via
     #   array-record
     #   clu
-    #   etils
     #   orbax-checkpoint
     #   tfds-nightly
 flatbuffers==23.5.26