skein launcher support large pex

cluster_pack/skein/skein_config_builder.py

@@ -83,12 +83,15 @@ def build(
     if not package_path:
         package_path, _ = uploader.upload_env()
 
+    python_env_descriptor = packaging.get_pyenv_usage_from_archive(package_path)
+
     script = _get_script(
-        package_path,
+        python_env_descriptor,
         module_name,
-        args)
+        args
+    )
 
-    files = _get_files(package_path, additional_files, tmp_dir)
+    files = _get_files(python_env_descriptor, additional_files, tmp_dir)
 
     env = {"SKEIN_CONFIG": "./.skein",
            "GIT_PYTHON_REFRESH": "quiet"}
@@ -118,17 +121,14 @@ def build(
 
 
 def _get_script(
-        package_path: str,
+        python_env_descriptor: packaging.PythonEnvDescription,
         module_name: str,
         args: List[Any] = []
 ) -> str:
-    python_bin = f"./{os.path.basename(package_path)}" if package_path.endswith(
-        '.pex') else f"./{os.path.basename(package_path)}/bin/python"
-
     launch_options = "-m" if not module_name.endswith(".py") else ""
     launch_args = " ".join(args)
 
-    cmd = f"{python_bin} {launch_options} {module_name} {launch_args}"
+    cmd = f"{python_env_descriptor.interpreter_cmd} {launch_options} {module_name} {launch_args}"
 
     script = f'''
                 export PEX_ROOT="./.pex"
@@ -141,17 +141,16 @@ def _get_script(
 
 
 def _get_files(
-        package_path: str,
+        python_env_descriptor: packaging.PythonEnvDescription,
         additional_files: Optional[List[str]] = None,
         tmp_dir: str = packaging._get_tmp_dir()
 ) -> Dict[str, str]:
-
-    files_to_upload = [package_path]
+    dict_files_to_upload = {
+        python_env_descriptor.dest_path: python_env_descriptor.path_to_archive
+    }
     if additional_files:
-        files_to_upload = files_to_upload + additional_files
-
-    dict_files_to_upload = {os.path.basename(path): path
-                            for path in files_to_upload}
+        for additional_file in additional_files:
+            dict_files_to_upload[os.path.basename(additional_file)] = additional_file
 
     editable_requirements = packaging.get_editable_requirements()
 

tests/test_packaging.py

@@ -169,13 +169,11 @@ def test_pack_in_pex(pyarrow_version, expectation):
     if sys.version_info.minor in {8, 9} and pyarrow_version == "0.13.0":
         return
     with tempfile.TemporaryDirectory() as tempdir:
-        requirements = ["tensorflow", f"pyarrow=={pyarrow_version}",
-                        "urllib3<2.0"]
-        packaging.pack_in_pex(
-            requirements,
-            f"{tempdir}/out.pex",
-            # make isolated pex from current pytest virtual env
-            pex_inherit_path="false")
+        requirements = [
+            "protobuf==3.19.6", "tensorflow==2.5.2",
+            "tensorboard==2.10.1", f"pyarrow=={pyarrow_version}"
+        ]
+        packaging.pack_in_pex(requirements, f"{tempdir}/out.pex", pex_inherit_path="false")
         assert os.path.exists(f"{tempdir}/out.pex")
         with expectation:
             print(subprocess.check_output([

tests/test_uploader.py

@@ -356,7 +356,7 @@ def test_format_pex_requirements():
             pex_inherit_path="false")
         pex_info = PexInfo.from_pex(f"{tempdir}/out.pex")
         cleaned_requirements = uploader._format_pex_requirements(pex_info)
-        pip_version = 'pip==21.3.1' if sys.version_info.minor == 6 else 'pip==23.1.2'
+        pip_version = 'pip==21.3.1' if sys.version_info.minor == 6 else 'pip==23.3'
         assert [pip_version, 'pipdeptree==2.0.0', 'six==1.15.0'] == cleaned_requirements