Merge pull request #127 from criteo/skein-allow-large-pex

Allow large PEX while submitting skein app

cluster_pack/skein/skein_config_builder.py

@@ -66,7 +66,8 @@ def build(
         package_path: Optional[str] = None,
         additional_files: Optional[List[str]] = None,
         tmp_dir: str = packaging._get_tmp_dir(),
-        process_logs: Callable[[str], Any] = None
+        process_logs: Callable[[str], Any] = None,
+        allow_large_pex: bool = False
 ) -> SkeinConfig:
     """Build the skein config for a module to execute
 
@@ -78,10 +79,13 @@ def build(
     :param tmp_dir: a temp dir for local files
     :param process_logs: hook with the local log path as a parameter,
                          can be used to uplaod the logs somewhere
+    :param allow_large_pex: Creates a non-executable pex that will need to be unzipped to circumvent
+                            python's limitation with zips > 2Gb. The file will need to be unzipped
+                            and the entry point will be <output>/__main__.py
     :return: SkeinConfig
     """
     if not package_path:
-        package_path, _ = uploader.upload_env()
+        package_path, _ = uploader.upload_env(allow_large_pex=allow_large_pex)
 
     python_env_descriptor = packaging.get_pyenv_usage_from_archive(package_path)
 

cluster_pack/skein/skein_launcher.py

@@ -25,7 +25,8 @@ def submit(skein_client: skein.Client,
            acquire_map_reduce_delegation_token: bool = False,
            pre_script_hook: Optional[str] = None,
            max_attempts: int = 1, max_restarts: int = 0,
-           process_logs: Callable[[str], Any] = None) -> str:
+           process_logs: Callable[[str], Any] = None,
+           allow_large_pex: bool = False) -> str:
     """Execute a python module in a skein container
 
     :param skein_client: skein.Client to use
@@ -50,6 +51,9 @@ def submit(skein_client: skein.Client,
     :param max_restarts: maximum number of restarts allowed for the service
     :param process_logs: hook with the local log path as a parameter,
                          can be used to uplaod the logs somewhere
+    :param allow_large_pex: Creates a non-executable pex that will need to be unzipped to circumvent
+                            python's limitation with zips > 2Gb. The file will need to be unzipped
+                            and the entry point will be <output>/__main__.py
     :return: SkeinConfig
     """
     with tempfile.TemporaryDirectory() as tmp_dir:
@@ -59,7 +63,8 @@ def submit(skein_client: skein.Client,
             package_path=package_path,
             additional_files=additional_files,
             tmp_dir=tmp_dir,
-            process_logs=process_logs)
+            process_logs=process_logs,
+            allow_large_pex=allow_large_pex)
 
         return _submit(
             skein_client, skein_config,