Adding pylint checks and code_style.sh to fix lint issues. Also addin…

…g files modified by code_style.sh (#114)

* add pylint files

* change files with code_style.sh

.github/workflows/CPUTests.yml

@@ -0,0 +1,46 @@
+name: Linter
+
+on:
+  push:
+    branches:
+      - '**'
+
+jobs:
+  cpu:
+    name: "CPU tests"
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        os: [ubuntu-20.04]
+        python-version: ['3.10']
+    steps:
+    - uses: actions/checkout@v3
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v4
+      with:
+        python-version: ${{ matrix.python-version }}
+    - name: Install Dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install pylint pyink pytype==2024.2.27
+    # - name: Typecheck the code with pytype
+    #   run: |
+    #     pytype --jobs auto --disable import-error src/maxdiffusion/
+    - name: Analysing the code with pylint in maxdiffusion/
+      run: |
+         pylint --fail-under=7 src/maxdiffusion/  && \
+         echo 'maxdiffusion PyLint check successful' || { echo \
+         'PyLint check has failed. Please run bash code_style.sh to fix issues'; exit 20; }
+    - name: Analysing the code with pylint in end_to_end/
+      run: |
+         pylint --fail-under=7 end_to_end/ && \
+         echo 'PyLint check on pend_to_end/ is successful' || { echo \
+         'PyLint check has failed. Please run bash code_style.sh to fix issues'; exit 20; }
+    - name: Analysing the code with pyink in maxdiffusion/
+      run: |
+        pyink src/maxdiffusion --check --diff --color --pyink-indentation=2 --line-length=125
+    - name: Analysing the code with pyink in end_to_end/
+      run: |
+        pyink end_to_end --check --diff --color --pyink-indentation=2 --line-length=125
+
+

code_style.sh

@@ -0,0 +1,40 @@
+# Copyright 2023 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.
+
+# Clean up Python codes using Pylint & Pyink
+# Googlers: please run `sudo apt install pipx; pipx install pylint --force; pipx install pyink==23.10.0` in advance
+
+set -e # Exit immediately if any command fails
+
+FOLDERS_TO_FORMAT=("src/maxdiffusion" "end_to_end/tpu")
+LINE_LENGTH=$(grep -E "^max-line-length=" pylintrc | cut -d '=' -f 2)
+
+# Check for --check flag
+CHECK_ONLY_PYINK_FLAGS=""
+if [[ "$1" == "--check" ]]; then
+  CHECK_ONLY_PYINK_FLAGS="--check --diff --color"
+fi
+
+for folder in "${FOLDERS_TO_FORMAT[@]}"
+do
+  pyink "$folder" ${CHECK_ONLY_PYINK_FLAGS} --pyink-indentation=2 --line-length=${LINE_LENGTH}
+done
+
+for folder in "${FOLDERS_TO_FORMAT[@]}"
+do
+  # pylint doesn't change files, only reports errors.
+  pylint "./$folder"
+done
+
+echo "Successfully clean up all codes."

end_to_end/tpu/eval_assert.py

@@ -20,9 +20,10 @@
 from typing import Sequence
 import json
 
+
 def get_last_n_data(metrics_file, target, n=10):
   last_n_data = []
-  with open(metrics_file, 'r', encoding='utf8') as file:
+  with open(metrics_file, "r", encoding="utf8") as file:
     lines = file.readlines()
     for line in lines[::-1]:
       metrics = json.loads(line)
@@ -35,24 +36,24 @@ def get_last_n_data(metrics_file, target, n=10):
 
 def test_final_loss(metrics_file, target_loss):
   target_loss = float(target_loss)
-  with open(metrics_file, 'r', encoding='utf8') as _:
+  with open(metrics_file, "r", encoding="utf8") as _:
     use_last_n_data = 10
-    last_n_data = get_last_n_data(metrics_file, 'learning/loss', use_last_n_data)
+    last_n_data = get_last_n_data(metrics_file, "learning/loss", use_last_n_data)
     avg_last_n_data = sum(last_n_data) / len(last_n_data)
     print(f"Mean of last {len(last_n_data)} losses is {avg_last_n_data}")
     print(f"Target loss is {target_loss}")
     assert avg_last_n_data < target_loss
-    print('Final loss test passed.')
+    print("Final loss test passed.")
 
 
 def main(argv: Sequence[str]) -> None:
 
   _, test_scenario, *test_vars = argv
 
-  if test_scenario == 'final_loss':
+  if test_scenario == "final_loss":
     test_final_loss(*test_vars)
   else:
-     raise ValueError(f"Unrecognized test_scenario {test_scenario}")
+    raise ValueError(f"Unrecognized test_scenario {test_scenario}")
 
 
 if __name__ == "__main__":