Skip to content

Commit

Permalink
Merge pull request #118 from criteo/check_platform_and_version
Browse files Browse the repository at this point in the history 
check version and platform when rebuilding pex
  • Loading branch information
@ax-vivien
ax-vivien authored Dec 5, 2023
2 parents 83df60d + 1915730 commit e536206
Show file tree
Hide file tree
Showing 2 changed files with 135 additions and 23 deletions.
40 changes: 37 additions & 3 deletions cluster_pack/uploader.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
import os
import sys
import pathlib
import platform
import tempfile
from typing import (
Tuple,
Expand All @@ -24,6 +25,10 @@

_logger = logging.getLogger(__name__)

PACKAGE_INSTALLED_KEY = "package_installed"
PLATFORM_KEY = "platform"
PYTHON_VERSION_KEY = "python_version"


def _get_archive_metadata_path(package_path: str) -> str:
url = parse.urlparse(package_path)
Expand All @@ -41,24 +46,53 @@ def _is_archive_up_to_date(package_path: str,
_logger.debug(f'metadata for archive {package_path} does not exist')
return False
with resolved_fs.open(archive_meta_data, "rb") as fd:
packages_installed = json.loads(fd.read())
return sorted(packages_installed) == sorted(current_packages_list)
metadata_dict = json.loads(fd.read())
if not isinstance(metadata_dict, dict):
_logger.debug('metadata exists but was built with old format')
return False

current_platform, current_python_version = get_platform_and_python_version()
packages_installed = metadata_dict.get(PACKAGE_INSTALLED_KEY, [])
platform = metadata_dict.get(PLATFORM_KEY, "")
python_version = metadata_dict.get(PYTHON_VERSION_KEY, "")
return (sorted(packages_installed) == sorted(current_packages_list)
and platform == current_platform
and python_version == current_python_version)


def _dump_archive_metadata(package_path: str,
current_packages_list: List[str],
resolved_fs: Any = None
) -> None:
archive_meta_data = _get_archive_metadata_path(package_path)
metadata_dict = build_metadata_dict(current_packages_list)
with tempfile.TemporaryDirectory() as tempdir:
tempfile_path = os.path.join(tempdir, "metadata.json")
with open(tempfile_path, "w") as fd:
fd.write(json.dumps(current_packages_list, sort_keys=True, indent=4))
fd.write(json.dumps(metadata_dict, indent=4))
if resolved_fs.exists(archive_meta_data):
resolved_fs.rm(archive_meta_data)
resolved_fs.put(tempfile_path, archive_meta_data)


def build_metadata_dict(current_packages_list: List[str]) -> Dict:
cur_platform, python_version = get_platform_and_python_version()
metadata_dict = {
PACKAGE_INSTALLED_KEY: current_packages_list,
PLATFORM_KEY: cur_platform,
PYTHON_VERSION_KEY: python_version
}
return metadata_dict


def get_platform_and_python_version() -> Tuple[str, str]:
system, _node, release, _version, _machine, _processor = platform.uname()
current_platform_str = f"{system}-{release}"
python_version = sys.version_info
python_version_str = f"{python_version.major}.{python_version.minor}.{python_version.micro}"
return current_platform_str, python_version_str


def upload_zip(
zip_file: str,
package_path: str = None,
Expand Down
118 changes: 98 additions & 20 deletions tests/test_uploader.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,17 +34,61 @@ def test_update_no_metadata():


@pytest.mark.parametrize("current_packages, metadata_packages, expected", [
pytest.param(["a==2.0", "b==1.0"], ["a==2.0", "b==1.0"], True),
pytest.param(["a==2.0", "b==1.0"], ["a==1.0", "b==1.0"], False),
pytest.param(["a==2.0", "b==1.0"], ["a==2.0"], False),
pytest.param(["a==2.0"], ["a==2.0", "b==1.0"], False),
pytest.param([], ["a==2.0", "b==1.0"], False),
pytest.param(["a==2.0"], ["c==1.0"], False),
pytest.param([], [], True),
pytest.param(["a==2.0", "b==1.0"],
{'package_installed': ["a==2.0", "b==1.0"],
'platform': 'fake_system-fake_release',
'python_version': '3.9.10'},
True),
pytest.param(["a==2.0", "b==1.0"],
{'package_installed': ["a==1.0", "b==1.0"],
'platform': 'fake_system-fake_release',
'python_version': '3.9.10'},
False),
pytest.param(["a==2.0", "b==1.0"],
{'package_installed': ["a==2.0"],
'platform': 'fake_system-fake_release',
'python_version': '3.9.10'},
False),
pytest.param(["a==2.0"],
{'package_installed': ["a==2.0", "b==1.0"],
'platform': 'fake_system-fake_release',
'python_version': '3.9.10'},
False),
pytest.param([],
{'package_installed': ["a==2.0", "b==1.0"],
'platform': 'fake_system-fake_release',
'python_version': '3.9.10'},
False),
pytest.param(["a==2.0"],
{'package_installed': ["c==1.0"],
'platform': 'fake_system-fake_release',
'python_version': '3.9.10'},
False),
pytest.param([],
{'package_installed': [],
'platform': 'fake_system-fake_release',
'python_version': '3.9.10'},
True),
pytest.param(["a==2.0", "b==1.0"],
["a==2.0", "b==1.0"],
False),
pytest.param(["a==2.0", "b==1.0"],
{'package_installed': ["a==2.0", "b==1.0"],
'platform': 'fake_system-fake_release2',
'python_version': '3.9.10'},
False),
pytest.param(["a==2.0", "b==1.0"],
{'package_installed': ["a==2.0", "b==1.0"],
'platform': 'fake_system-fake_release',
'python_version': '3.6.8'},
False),
])
def test_update_version_comparaison(current_packages, metadata_packages,
expected):

@mock.patch(f'{MODULE_TO_TEST}.platform.uname')
@mock.patch(f'{MODULE_TO_TEST}.sys')
def test_update_version_comparaison(sys_mock, uname_mock,
current_packages, metadata_packages, expected):
uname_mock.return_value = build_uname_mock()
sys_mock.version_info = build_python_version_mock()
map_is_exist = {MYARCHIVE_FILENAME: True,
MYARCHIVE_METADATA: True}

Expand Down Expand Up @@ -74,21 +118,40 @@ def __eq__(self, other):
}"""


def test_dump_metadata():
def build_python_version_mock():
python_version_mock = mock.Mock()
python_version_mock.major = 3
python_version_mock.minor = 9
python_version_mock.micro = 10
return python_version_mock


def build_uname_mock():
return "fake_system", "", "fake_release", "", "", ""


@mock.patch(f'{MODULE_TO_TEST}.platform.uname')
@mock.patch(f'{MODULE_TO_TEST}.sys')
def test_dump_metadata(sys_mock, uname_mock):
uname_mock.return_value = build_uname_mock()
sys_mock.version_info = build_python_version_mock()
mock_fs = mock.Mock()
mock_fs.rm.return_value = True
mock_fs.exists.return_value = True
mock_open = mock.mock_open()
with mock.patch.object(mock_fs, 'open', mock_open):
mock_fs.exists.return_value = True
packages = {"a": "1.0", "b": "2.0"}
packages = ["a==1.0", "b==2.0"]
uploader._dump_archive_metadata(
MYARCHIVE_FILENAME,
packages,
filesystem.EnhancedFileSystem(mock_fs))
# Check previous file has been deleted
mock_fs.rm.assert_called_once_with(MYARCHIVE_METADATA)
mock_open().write.assert_called_once_with(b'{\n "a": "1.0",\n "b": "2.0"\n}')
mock_open().write.assert_called_once_with(
b'{\n "package_installed": [\n "a==1.0",\n "b==2.0"\n ],'
+ b'\n "platform": "fake_system-fake_release",\n "python_version": "3.9.10"\n}'
)


def test_upload_env():
Expand Down Expand Up @@ -238,17 +301,26 @@ def test_upload_spec_hdfs(
assert result_path == "hdfs:///user/testuser/envs/myenv.pex"


def test_upload_spec_local_fs():
@mock.patch(f'{MODULE_TO_TEST}.platform.uname')
@mock.patch(f'{MODULE_TO_TEST}.sys')
def test_upload_spec_local_fs(sys_mock, uname_mock):
uname_mock.return_value = build_uname_mock()
sys_mock.version_info = build_python_version_mock()
spec_file = os.path.join(os.path.dirname(__file__), "resources", "requirements.txt")
with tempfile.TemporaryDirectory() as tempdir:
result_path = cluster_pack.upload_spec(spec_file, f"{tempdir}/package.pex")
assert os.path.exists(result_path)
_check_metadata(
f"{tempdir}/package.json",
["5a5f33b106aad8584345f5a0044a4188ce78b3f4"])
{'package_installed': ['5a5f33b106aad8584345f5a0044a4188ce78b3f4'],
'platform': 'fake_system-fake_release', 'python_version': '3.9.10'})


def test_upload_spec_unique_name():
@mock.patch(f'{MODULE_TO_TEST}.platform.uname')
@mock.patch(f'{MODULE_TO_TEST}.sys')
def test_upload_spec_unique_name(sys_mock, uname_mock):
uname_mock.return_value = build_uname_mock()
sys_mock.version_info = build_python_version_mock()
with tempfile.TemporaryDirectory() as tempdir:
spec_file = f"{tempdir}/myproject/requirements.txt"
_write_spec_file(spec_file, ["cloudpickle==1.4.1"])
Expand All @@ -259,7 +331,8 @@ def test_upload_spec_unique_name():
assert result_path == f"{tempdir}/cluster_pack_myproject.pex"
_check_metadata(
f"{tempdir}/cluster_pack_myproject.json",
["b8721a3c125d3f7edfa27d7b13236e696f652a16"])
{'package_installed': ["b8721a3c125d3f7edfa27d7b13236e696f652a16"],
'platform': 'fake_system-fake_release', 'python_version': '3.9.10'})


@mock.patch(f"{MODULE_TO_TEST}.packaging.pack_spec_in_pex")
Expand All @@ -272,7 +345,6 @@ def test_upload_spec_local_fs_use_cache(mock_pack_spec_in_pex):
mock_pack_spec_in_pex.return_value = pex_file
with open(pex_file, "w"):
pass

result_path = cluster_pack.upload_spec(spec_file, pex_file)
result_path1 = cluster_pack.upload_spec(spec_file, pex_file)

Expand All @@ -281,9 +353,14 @@ def test_upload_spec_local_fs_use_cache(mock_pack_spec_in_pex):
assert result_path == result_path1 == pex_file


@mock.patch(f'{MODULE_TO_TEST}.platform.uname')
@mock.patch(f'{MODULE_TO_TEST}.sys')
@mock.patch(f"{MODULE_TO_TEST}.packaging.pack_spec_in_pex")
def test_upload_spec_local_fs_changed_reqs(mock_pack_spec_in_pex):
def test_upload_spec_local_fs_changed_reqs(mock_pack_spec_in_pex, sys_mock, uname_mock):
mock_pack_spec_in_pex.return_value = "/tmp/tmp.pex"
uname_mock.return_value = build_uname_mock()
sys_mock.version_info = build_python_version_mock()

with tempfile.TemporaryDirectory() as tempdir:
spec_file = f"{tempdir}/myproject/requirements.txt"
_write_spec_file(spec_file, ["cloudpickle==1.4.1"])
Expand All @@ -308,7 +385,8 @@ def test_upload_spec_local_fs_changed_reqs(mock_pack_spec_in_pex):
assert os.path.exists(result_path1)
_check_metadata(
f"{tempdir}/package.json",
["0fd17ced922a2387fa660fb0cb78e1c77fbe3349"])
{'package_installed': ['0fd17ced922a2387fa660fb0cb78e1c77fbe3349'],
'platform': 'fake_system-fake_release', 'python_version': '3.9.10'})


def test__handle_packages_use_local_wheel():
Expand Down

0 comments on commit e536206

Please sign in to comment.