refactor: run post_setup script on init (#87)

made changes to that post_setup script is executed in __init__.py of multiworld
error handlers were while executing the main method.
with the first execution, the patch is applied and the user is prompted to re-run the script to have the patched version of torch

Co-authored-by: Rares Gaia <[email protected]>

multiworld/__init__.py

@@ -22,8 +22,12 @@
 
 from multiworld.version import VERSION as __version__  # noqa: F401
 
+from . import post_setup
+
 logging.basicConfig(
     level=getattr(logging, os.getenv("M8D_LOG_LEVEL", "WARNING")),
     format="%(asctime)s | %(filename)s:%(lineno)d | %(levelname)s | %(threadName)s | %(funcName)s | %(message)s",
     handlers=[logging.StreamHandler(sys.stdout)],
 )
+
+post_setup.configure_once()

multiworld/init.txt

@@ -0,0 +1 @@
+false

multiworld/post_setup.py

@@ -14,36 +14,39 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-import argparse
 import os
 import pathlib
 import shutil
 import site
+import sys
 
-import torch
 
+def configure_once():
+    package_name = __name__.split(".")[0]
+    path_to_sitepackages = site.getsitepackages()[0]
 
-def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "patchfile", nargs="?", default=None, help="Path to the patch file"
-    )
-    args = parser.parse_args()
+    init_file_path = os.path.join(path_to_sitepackages, package_name, "init.txt")
 
-    path_to_sitepackages = site.getsitepackages()[0]
+    with open(init_file_path, "r") as file:
+        patch_applied = file.read()
+
+    if patch_applied == "true":
+        return
+
+    print(f"Configuring {package_name} for the first time. This is one time task.")
 
-    if args.patchfile:
-        patchfile = args.patchfile
-    else:
-        torch_version = torch.__version__.split("+")[
-            0
-        ]  # torch version is in "2.2.1+cu121" format
-        patchfile = os.path.join(
-            path_to_sitepackages,
-            "multiworld",
-            "patch",
-            "pytorch-v" + torch_version + ".patch",
-        )
+    import torch
+
+    torch_version = torch.__version__.split("+")[
+        0
+    ]  # torch version is in "2.2.1+cu121" format
+
+    patchfile = os.path.join(
+        path_to_sitepackages,
+        package_name,
+        "patch",
+        "pytorch-v" + torch_version + ".patch",
+    )
 
     patch_basename = os.path.basename(patchfile)
 
@@ -52,10 +55,13 @@ def main():
 
     os.chdir(path_to_sitepackages)
 
-    os.system(f"patch -p1 < {patch_basename}")
+    os.system(f"patch -p1 < {patch_basename} > /dev/null")
     p = pathlib.Path(patch_basename)
     p.unlink()
 
+    with open(init_file_path, "w") as file:
+        file.write("true")
 
-if __name__ == "__main__":
-    main()
+    sys.exit(
+        f"This one-time configuration for {package_name} is completed.\nYou can run your script without any interruption from now on."
+    )

pyproject.toml

@@ -13,9 +13,6 @@ dependencies = [
     "torch == 2.4.0",
 ]
 
-[project.scripts]
-m8d-post-setup = "multiworld.post_setup:main"
-
 [Project.optional-dependencies]
 dev = [
     "black",
@@ -30,7 +27,7 @@ dev = [
 packages=["multiworld"]
 
 [tool.setuptools.package-data]
-"multiworld" = ["patch/*.patch"]
+"multiworld" = ["patch/*.patch", "init.txt"]
 
 [tool.setuptools.dynamic]
 version = {attr = "multiworld.__version__"}