Torchelastic bring back setup.py

Summary: Restore setup.py that requires for building docs.

Reviewed By: kiukchung

Differential Revision: D27974755

fbshipit-source-id: 2c2ecd9f9b4a9cd0f3415991ff05270ebb6ef195

Author: aivanou

__init__.py

@@ -0,0 +1,7 @@
+#!/usr/bin/env python3
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.

setup.py

@@ -7,7 +7,54 @@
 # LICENSE file in the root directory of this source tree.
 
 
+import sys
+
+from setuptools import find_packages, setup
+
+
+def get_version():
+    return "0.2.3.dev0"
+
+
 if __name__ == "__main__":
-    raise RuntimeError(
-        "Torchelastic got merged with pytorch. Use pytorch: https://pytorch.org/"
+    if sys.version_info < (3, 8):
+        sys.exit("python >= 3.8 required for torchelastic")
+
+    with open("README.md", encoding="utf8") as f:
+        readme = f.read()
+
+    with open("requirements.txt") as f:
+        reqs = f.read()
+
+    version = get_version()
+    print("-- Building version: " + version)
+
+    setup(
+        # Metadata
+        name="torchelastic",
+        version=version,
+        author="PyTorch Elastic Devs",
+        author_email="[email protected]",
+        description="PyTorch Elastic Training",
+        long_description=readme,
+        long_description_content_type="text/markdown",
+        url="https://github.com/pytorch/elastic",
+        license="BSD-3",
+        keywords=["pytorch", "machine learning", "elastic", "distributed"],
+        python_requires=">=3.8",
+        install_requires=reqs.strip().split("\n"),
+        include_package_data=True,
+        packages=find_packages(exclude=("*.test", "aws*", "*.fb")),
+        test_suite="torchelastic.tsm.test.suites.unittests",
+        # PyPI package information.
+        classifiers=[
+            "Development Status :: 4 - Beta",
+            "Intended Audience :: Developers",
+            "Intended Audience :: Science/Research",
+            "License :: OSI Approved :: BSD License",
+            "Programming Language :: Python :: 3",
+            "Programming Language :: Python :: 3.8",
+            "Topic :: System :: Distributed Computing",
+            "Topic :: Scientific/Engineering :: Artificial Intelligence",
+        ],
     )

torchelastic/__init__.py

@@ -0,0 +1,7 @@
+#!/usr/bin/env python3
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.

torchelastic/tsm/test/suites.py

@@ -0,0 +1,50 @@
+#!/usr/bin/env python3
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+import random
+import unittest
+from itertools import chain
+
+
+def _circleci_parallelism(suite):
+    """Allow for parallelism in CircleCI for speedier tests.."""
+    if int(os.environ.get("CIRCLE_NODE_TOTAL", 0)) <= 1:
+        # either not running on circleci, or we're not using parallelism.
+        return suite
+    # tests are automatically sorted by discover, so we will get the same ordering
+    # on all hosts.
+    total = int(os.environ["CIRCLE_NODE_TOTAL"])
+    index = int(os.environ["CIRCLE_NODE_INDEX"])
+
+    # right now each test is corresponds to a /file/. Certain files are slower than
+    # others, so we want to flatten it
+    tests = [testfile._tests for testfile in suite._tests]
+    tests = list(chain.from_iterable(tests))
+    random.Random(42).shuffle(tests)
+    tests = [t for i, t in enumerate(tests) if i % total == index]
+    return unittest.TestSuite(tests)
+
+
+def unittests():
+    """
+    Short tests.
+
+    Runs on CircleCI on every commit. Returns everything in the tests root directory.
+    """
+    test_loader = unittest.TestLoader()
+    test_suite = test_loader.discover(
+        "torchelastic/tsm", pattern="*_test.py", top_level_dir="."
+    )
+    test_suite = _circleci_parallelism(test_suite)
+    return test_suite
+
+
+if __name__ == "__main__":
+    runner = unittest.TextTestRunner()
+    runner.run(unittests())