diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index a776632..bdf57bf 100644 --- a/.github/workflows/main.yml +++ b/.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 diff --git a/README.md b/README.md index 8108941..b5f4943 100644 --- a/README.md +++ b/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 diff --git a/cluster_pack/packaging.py b/cluster_pack/packaging.py index 3f0d182..22f3432 100644 --- a/cluster_pack/packaging.py +++ b/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)}" diff --git a/cluster_pack/skein/skein_launcher.py b/cluster_pack/skein/skein_launcher.py index 1195bcf..a54f8e2 100644 --- a/cluster_pack/skein/skein_launcher.py +++ b/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: diff --git a/cluster_pack/uploader.py b/cluster_pack/uploader.py index bedbcc1..b7502cc 100644 --- a/cluster_pack/uploader.py +++ b/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}" diff --git a/tests/test_packaging.py b/tests/test_packaging.py index c111968..54534d0 100644 --- a/tests/test_packaging.py +++ b/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"), ] diff --git a/tests/test_uploader.py b/tests/test_uploader.py index f6e99c4..4ca8496 100644 --- a/tests/test_uploader.py +++ b/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