Remove main function entrypoint in ModelBuilder dependency manager. (aws#5058)

* Remove main function entrypoint in ModelBuilder dependency manager.

* Remove main function entrypoint in ModelBuilder dependency manager.

---------

Co-authored-by: Joseph Zhang <[email protected]>

Author: 2 people

src/sagemaker/serve/detector/dependency_manager.py

@@ -34,22 +34,34 @@ def capture_dependencies(dependencies: dict, work_dir: Path, capture_all: bool =
     """Placeholder docstring"""
     path = work_dir.joinpath("requirements.txt")
     if "auto" in dependencies and dependencies["auto"]:
+        import site
+
+        pkl_path = work_dir.joinpath(PKL_FILE_NAME)
+        dest_path = path
+        site_packages_dir = site.getsitepackages()[0]
+        pickle_command_dir = "/sagemaker/serve/detector"
+
         command = [
             sys.executable,
-            Path(__file__).parent.joinpath("pickle_dependencies.py"),
-            "--pkl_path",
-            work_dir.joinpath(PKL_FILE_NAME),
-            "--dest",
-            path,
+            "-c",
         ]
 
         if capture_all:
-            command.append("--capture_all")
+            command.append(
+                f"from pickle_dependencies import get_all_requirements;"
+                f'get_all_requirements("{dest_path}")'
+            )
+        else:
+            command.append(
+                f"from pickle_dependencies import get_requirements_for_pkl_file;"
+                f'get_requirements_for_pkl_file("{pkl_path}", "{dest_path}")'
+            )
 
         subprocess.run(
             command,
             env={"SETUPTOOLS_USE_DISTUTILS": "stdlib"},
             check=True,
+            cwd=site_packages_dir + pickle_command_dir,
         )
 
         with open(path, "r") as f:

src/sagemaker/serve/detector/pickle_dependencies.py

@@ -3,7 +3,6 @@
 from __future__ import absolute_import
 from pathlib import Path
 from typing import List
-import argparse
 import email.parser
 import email.policy
 import json
@@ -129,32 +128,3 @@ def get_all_requirements(dest: Path):
             version = package_info.get("version")
 
             out.write(f"{name}=={version}\n")
-
-
-def parse_args():
-    """Placeholder docstring"""
-    parser = argparse.ArgumentParser(
-        prog="pkl_requirements", description="Generates a requirements.txt for a cloudpickle file"
-    )
-    parser.add_argument("--pkl_path", required=True, help="path of the pkl file")
-    parser.add_argument("--dest", required=True, help="path of the destination requirements.txt")
-    parser.add_argument(
-        "--capture_all",
-        action="store_true",
-        help="capture all dependencies in current environment",
-    )
-    args = parser.parse_args()
-    return (Path(args.pkl_path), Path(args.dest), args.capture_all)
-
-
-def main():
-    """Placeholder docstring"""
-    pkl_path, dest, capture_all = parse_args()
-    if capture_all:
-        get_all_requirements(dest)
-    else:
-        get_requirements_for_pkl_file(pkl_path, dest)
-
-
-if __name__ == "__main__":
-    main()