From c8d69c8b4f19112f6c69af42421ab0bb6d8b7a59 Mon Sep 17 00:00:00 2001
From: Joseph Zhang <cjz@amazon.com>
Date: Mon, 24 Feb 2025 17:00:40 -0800
Subject: [PATCH 1/2] Remove main function entrypoint in ModelBuilder
 dependency manager.

---
 .../serve/detector/dependency_manager.py      | 24 +++++++++++----
 .../serve/detector/pickle_dependencies.py     | 29 -------------------
 2 files changed, 18 insertions(+), 35 deletions(-)

diff --git a/src/sagemaker/serve/detector/dependency_manager.py b/src/sagemaker/serve/detector/dependency_manager.py
index e72a84da30..8ff37c9185 100644
--- a/src/sagemaker/serve/detector/dependency_manager.py
+++ b/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:
diff --git a/src/sagemaker/serve/detector/pickle_dependencies.py b/src/sagemaker/serve/detector/pickle_dependencies.py
index 5a1cd43869..84fefba6bf 100644
--- a/src/sagemaker/serve/detector/pickle_dependencies.py
+++ b/src/sagemaker/serve/detector/pickle_dependencies.py
@@ -129,32 +129,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()

From b9e2517361b3ad5bf1032635c10d2c8dd5cc30fe Mon Sep 17 00:00:00 2001
From: Joseph Zhang <cjz@amazon.com>
Date: Mon, 24 Feb 2025 17:00:40 -0800
Subject: [PATCH 2/2] Remove main function entrypoint in ModelBuilder
 dependency manager.

---
 src/sagemaker/serve/detector/pickle_dependencies.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/sagemaker/serve/detector/pickle_dependencies.py b/src/sagemaker/serve/detector/pickle_dependencies.py
index 84fefba6bf..8f9da917fd 100644
--- a/src/sagemaker/serve/detector/pickle_dependencies.py
+++ b/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