Cherry pick release 0.3.17.post1

.github/workflows/main.yml

@@ -54,7 +54,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        version: [ { python: 3.9.15, spark: 3.5.7 }, { python: 3.11.8, spark: 3.5.7 } ]
+        version: [ { python: 3.9.15, spark: 3.5.8 }, { python: 3.11.8, spark: 3.5.8 } ]
 
     steps:
     - uses: actions/checkout@v2

README.md

@@ -30,6 +30,12 @@ $ pip install .
 
 cluster-pack supports Python ≥3.9.
 
+## Feature flags
+- C_PACK_USER: override the current user for HDFS path generation and Skein impersonation
+  - When set, this value is used instead of the system user (from `getpass.getuser()`)
+  - Useful for running jobs as a different user or in environments where the system user doesn't match the HDFS user
+  - If not set or empty, falls back to the current system user
+
 ## Features
 
 - Ships a package with all the dependencies from your current virtual environment

cluster_pack/packaging.py

@@ -1,5 +1,5 @@
 import getpass
-import imp
+import importlib.util
 import json
 import logging
 import os
@@ -53,6 +53,21 @@ def _get_tmp_dir() -> str:
     return tmp_dir
 
 
+def _get_current_user() -> str:
+    """
+    Get the current user name.
+
+    First checks the C_PACK_USER environment variable.
+    If not set or empty, falls back to getpass.getuser().
+
+    :return: username string
+    """
+    user = os.environ.get("C_PACK_USER", "").strip()
+    if user:
+        return user
+    return getpass.getuser()
+
+
 def zip_path(
     py_dir: str, include_base_name: bool = True, tmp_dir: str = _get_tmp_dir()
 ) -> str:
@@ -373,7 +388,7 @@ def get_non_editable_requirements(executable: str = sys.executable) -> Dict[str,
 
 
 def _build_package_path(name: str, extension: Optional[str]) -> str:
-    path = f"{get_default_fs()}/user/{getpass.getuser()}/envs/{name}"
+    path = f"{get_default_fs()}/user/{_get_current_user()}/envs/{name}"
     if extension is None:
         return path
     return f"{path}.{extension}"
@@ -489,9 +504,12 @@ def get_editable_requirements(
         else:
             for package_name in package_names:
                 try:
-                    _, path, _ = imp.find_module(package_name)
+                    spec = importlib.util.find_spec(package_name)
+                    if spec is None or spec.origin is None:
+                        raise ModuleNotFoundError(f"No module named '{package_name}'")
+                    path = os.path.dirname(spec.origin)
                     editable_requirements[os.path.basename(path)] = path
-                except ImportError:
+                except ModuleNotFoundError:
                     _logger.error(
                         f"Could not import package {package_name}"
                         f" repo exists={os.path.exists(package_name)}"

cluster_pack/skein/skein_launcher.py

@@ -10,6 +10,7 @@
 
 from cluster_pack.skein import skein_config_builder
 from cluster_pack import filesystem
+from cluster_pack.packaging import _get_current_user
 
 logger = logging.getLogger(__name__)
 
@@ -223,7 +224,7 @@ def _submit(
         spec.acquire_map_reduce_delegation_token = acquire_map_reduce_delegation_token
 
     # activate Impersonation only if user to run the job is not the current user (yarn issue)
-    if user and user != getpass.getuser():
+    if user and user != _get_current_user():
         spec.user = user
 
     if queue:

cluster_pack/uploader.py

@@ -271,7 +271,7 @@ def _upload_pex_file(
             info_to_upload = PexInfo.from_pex(pex_file_or_dir)
             if (
                 not force_upload
-                and info_from_storage.code_hash == info_to_upload.code_hash
+                and info_from_storage.pex_hash == info_to_upload.pex_hash
             ):
                 _logger.info(
                     f"skip upload of current {pex_file_or_dir}"

tests/test_packaging.py

@@ -22,7 +22,8 @@
     LARGE_PEX_CMD,
     resolve_zip_from_pex_dir,
     check_large_pex,
-    PexTooLargeError
+    PexTooLargeError,
+    _get_current_user
 )
 
 MODULE_TO_TEST = "cluster_pack.packaging"
@@ -44,6 +45,55 @@ def test_get_virtualenv_empty_returns_default():
         assert "default" == packaging.get_env_name(VARNAME)
 
 
+def test_get_current_user_with_env_variable():
+    """Test that C_PACK_USER environment variable is used when set."""
+    with mock.patch.dict("os.environ"):
+        # Test when C_PACK_USER is set
+        os.environ["C_PACK_USER"] = "custom_user"
+        assert packaging._get_current_user() == "custom_user"
+
+
+def test_get_current_user_with_empty_env_variable():
+    """Test that empty C_PACK_USER falls back to getpass.getuser()."""
+    with contextlib.ExitStack() as stack:
+        stack.enter_context(
+            mock.patch(f"{MODULE_TO_TEST}.getpass.getuser", return_value="system_user")
+        )
+        with mock.patch.dict("os.environ", clear=True):
+            # Test when C_PACK_USER is not set
+            assert packaging._get_current_user() == "system_user"
+
+            # Test when C_PACK_USER is empty string
+            os.environ["C_PACK_USER"] = ""
+            assert packaging._get_current_user() == "system_user"
+
+            # Test when C_PACK_USER is only whitespace
+            os.environ["C_PACK_USER"] = "   "
+            assert packaging._get_current_user() == "system_user"
+
+
+def test_get_current_user_strips_whitespace():
+    """Test that C_PACK_USER whitespace is stripped."""
+    with mock.patch.dict("os.environ"):
+        os.environ["C_PACK_USER"] = "  spaced_user  "
+        assert packaging._get_current_user() == "spaced_user"
+
+
+def test_build_package_path_uses_c_pack_user_env():
+    """Test that _build_package_path uses C_PACK_USER when set."""
+    with contextlib.ExitStack() as stack:
+        stack.enter_context(
+            mock.patch(f"{MODULE_TO_TEST}.get_default_fs", return_value="hdfs://")
+        )
+        with mock.patch.dict("os.environ"):
+            os.environ["C_PACK_USER"] = "env_user"
+
+            result = packaging._build_package_path("myenv", "pex")
+            expected = "hdfs:///user/env_user/envs/myenv.pex"
+
+            assert result == expected
+
+
 def test_get_empty_editable_requirements():
     with tempfile.TemporaryDirectory() as tempdir:
         _create_venv(tempdir)
@@ -243,18 +293,10 @@ def does_not_raise():
 
 
 def test_pack_in_pex():
-    if sys.version_info.minor == 9:
-        requirements = [
-            "protobuf==3.19.6",
-            "tensorflow==2.5.2",
-            "tensorboard==2.10.1",
-            "pyarrow==6.0.1",
-        ]
-    else:
-        requirements = [
-            "tensorflow",
-            "pyarrow"
-        ]
+    requirements = [
+        "numpy",
+        "pyarrow"
+    ]
     with tempfile.TemporaryDirectory() as tempdir:
         packaging.pack_in_pex(
             requirements, f"{tempdir}/out.pex", pex_inherit_path="false"
@@ -267,9 +309,9 @@ def test_pack_in_pex():
                         f"{tempdir}/out.pex",
                         "-c",
                         (
-                            """print("Start importing pyarrow and tensorflow..");"""
-                            """import pyarrow; import tensorflow;"""
-                            """print("Successfully imported pyarrow and tensorflow!")"""
+                            """print("Start importing pyarrow and numpy..");"""
+                            """import pyarrow; import numpy;"""
+                            """print("Successfully imported pyarrow and numpy!")"""
                         ),
                     ]
                 )
@@ -530,19 +572,19 @@ def test_gen_pyenvs_from_unknown_format():
     (False, "dummy/path/exe.pex", True, False, "dummy/path/exe.pex.zip"),
     (True, "dummy/path/exe.pex", False, False, "dummy/path/exe.pex"),
     (True, "dummy/path/exe.pex", True, False, "dummy/path/exe.pex.zip"),
-    (False, None, False, False, f"hdfs:///user/{getpass.getuser()}/envs/venv_exe.pex"),
-    (False, None, None, False, f"hdfs:///user/{getpass.getuser()}/envs/venv_exe.pex"),
+    (False, None, False, False, f"hdfs:///user/{_get_current_user()}/envs/venv_exe.pex"),
+    (False, None, None, False, f"hdfs:///user/{_get_current_user()}/envs/venv_exe.pex"),
     (
         False,
         None,
         True,
         False,
-        f"hdfs:///user/{getpass.getuser()}/envs/venv_exe.pex.zip",
+        f"hdfs:///user/{_get_current_user()}/envs/venv_exe.pex.zip",
     ),
-    (True, None, False, False, f"hdfs:///user/{getpass.getuser()}/envs/pex_exe.pex"),
-    (True, None, True, False, f"hdfs:///user/{getpass.getuser()}/envs/pex_exe.pex.zip"),
-    (True, None, None, False, f"hdfs:///user/{getpass.getuser()}/envs/pex_exe.pex"),
-    (True, None, None, True, f"hdfs:///user/{getpass.getuser()}/envs/pex_exe.pex.zip"),
+    (True, None, False, False, f"hdfs:///user/{_get_current_user()}/envs/pex_exe.pex"),
+    (True, None, True, False, f"hdfs:///user/{_get_current_user()}/envs/pex_exe.pex.zip"),
+    (True, None, None, False, f"hdfs:///user/{_get_current_user()}/envs/pex_exe.pex"),
+    (True, None, None, True, f"hdfs:///user/{_get_current_user()}/envs/pex_exe.pex.zip"),
 ]
 
 

tests/test_uploader.py

@@ -327,9 +327,9 @@ def test_upload_env_in_a_pex():
 
         def _from_pex(arg):
             if arg == f"{home_path}/myapp.pex":
-                return PexInfo({"code_hash": 1})
+                return PexInfo({"pex_hash": 1})
             else:
-                return PexInfo({"code_hash": 2})
+                return PexInfo({"pex_hash": 2})
 
         mock_pex_info.from_pex.side_effect = _from_pex