Adding a warning when user tries to load train.py in the old manner (as a script)

MaxText/tests/module_tests.py

@@ -0,0 +1,50 @@
+"""
+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
+
+     https://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.
+"""
+
+"""Tests for train.py with various configs"""
+
+import os
+import sys
+import subprocess
+import unittest
+from absl.testing import absltest
+
+from MaxText.globals import PKG_DIR
+
+
+class ModuleTests(unittest.TestCase):
+  """Tests train.py with various invocation methods"""
+
+  def test_get_informative_error(self):
+    command = [
+        sys.executable,  # use the same interpreter instance
+        "MaxText/train.py",
+        "MaxText/configs/base.yml",
+        "run_name=maxtext-module-test",
+        "base_output_directory=gs://does-not-exist",
+        "dataset_type=synthetic",
+    ]
+
+    with self.assertRaises(subprocess.CalledProcessError) as context:
+      subprocess.run(command, check=True, text=True, capture_output=True)
+
+    ex = context.exception
+    self.assertEqual(ex.returncode, 64)
+    self.assertIn("The MaxText API has changed", ex.stderr)
+
+
+if __name__ == "__main__":
+  absltest.main()

MaxText/train.py

@@ -17,8 +17,18 @@
 # pylint: disable=g-bad-todo, abstract-method, consider-using-with, ungrouped-imports
 """Training loop and Decoding of the model."""
 
-# Calling jax.device_count here prevents a "TPU platform already registered" error.
-# See github.com/google/maxtext/issues/20 for more
+import sys
+
+if not __package__:
+  print(
+      "Error: The MaxText API has changed. MaxText entry-points are now "
+      "invoked as modules, with syntax of the form "
+      "`python3 -m MaxText.module <args>`, e.g., "
+      "`python3 -m MaxText.train <args>`.",
+      file=sys.stderr,
+  )
+  # EX_USAGE
+  sys.exit(64)
 
 import datetime
 import os