|
| 1 | +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. |
| 2 | +"""Set up locally built adaptor wheels on a Deadline Cloud worker.""" |
| 3 | + |
| 4 | +from __future__ import annotations |
| 5 | + |
| 6 | +import argparse |
| 7 | +import os |
| 8 | +import re |
| 9 | +import subprocess |
| 10 | +import sys |
| 11 | +from collections.abc import Sequence |
| 12 | +from importlib.metadata import distributions |
| 13 | +from pathlib import Path |
| 14 | + |
| 15 | +_EXPECTED_WHEEL_PREFIXES = ( |
| 16 | + "openjd_adaptor_runtime-", |
| 17 | + "deadline-", |
| 18 | + "deadline_cloud_for_cinema_4d-", |
| 19 | +) |
| 20 | + |
| 21 | + |
| 22 | +def _find_wheels(wheels_dir: Path) -> list[Path]: |
| 23 | + wheels = sorted(wheels_dir.glob("*.whl")) |
| 24 | + selected: list[Path] = [] |
| 25 | + |
| 26 | + for prefix in _EXPECTED_WHEEL_PREFIXES: |
| 27 | + matches = [wheel for wheel in wheels if wheel.name.startswith(prefix)] |
| 28 | + if len(matches) != 1: |
| 29 | + raise RuntimeError( |
| 30 | + f"Expected exactly one wheel matching '{prefix}*.whl' in {wheels_dir}, " |
| 31 | + f"found {[wheel.name for wheel in matches]}" |
| 32 | + ) |
| 33 | + selected.append(matches[0]) |
| 34 | + |
| 35 | + if len(wheels) != len(selected): |
| 36 | + unexpected = sorted(set(wheels) - set(selected)) |
| 37 | + raise RuntimeError( |
| 38 | + "The adaptor wheels directory contains unexpected wheels: " |
| 39 | + + ", ".join(wheel.name for wheel in unexpected) |
| 40 | + ) |
| 41 | + |
| 42 | + return selected |
| 43 | + |
| 44 | + |
| 45 | +def _get_venv_paths(venv_dir: Path, adaptor_name: str) -> tuple[Path, Path, Path]: |
| 46 | + if os.name == "nt": |
| 47 | + bin_dir = venv_dir / "Scripts" |
| 48 | + return bin_dir, bin_dir / "python.exe", bin_dir / f"{adaptor_name}.exe" |
| 49 | + |
| 50 | + bin_dir = venv_dir / "bin" |
| 51 | + return bin_dir, bin_dir / "python", bin_dir / adaptor_name |
| 52 | + |
| 53 | + |
| 54 | +def _get_site_packages(venv_python: Path) -> Path: |
| 55 | + result = subprocess.run( |
| 56 | + [ |
| 57 | + str(venv_python), |
| 58 | + "-c", |
| 59 | + "import sysconfig; print(sysconfig.get_paths()['purelib'])", |
| 60 | + ], |
| 61 | + check=True, |
| 62 | + capture_output=True, |
| 63 | + text=True, |
| 64 | + ) |
| 65 | + return Path(result.stdout.strip()) |
| 66 | + |
| 67 | + |
| 68 | +def _get_distribution_version(site_packages: Path, distribution_name: str) -> str: |
| 69 | + canonical_name = re.sub(r"[-_.]+", "-", distribution_name).lower() |
| 70 | + for distribution in distributions(path=[str(site_packages)]): |
| 71 | + installed_name = distribution.metadata["Name"] |
| 72 | + if installed_name and re.sub(r"[-_.]+", "-", installed_name).lower() == canonical_name: |
| 73 | + return distribution.version |
| 74 | + raise RuntimeError(f"Could not find '{distribution_name}' in {site_packages}") |
| 75 | + |
| 76 | + |
| 77 | +def _install_wheels(venv_python: Path, wheels: list[Path]) -> None: |
| 78 | + # The active Conda environment supplies transitive dependencies. Development |
| 79 | + # wheels use generated versions that may not satisfy each other's release ranges. |
| 80 | + subprocess.run( |
| 81 | + [ |
| 82 | + str(venv_python), |
| 83 | + "-m", |
| 84 | + "pip", |
| 85 | + "install", |
| 86 | + "--disable-pip-version-check", |
| 87 | + "--force-reinstall", |
| 88 | + "--no-deps", |
| 89 | + *[str(wheel) for wheel in wheels], |
| 90 | + ], |
| 91 | + check=True, |
| 92 | + ) |
| 93 | + |
| 94 | + |
| 95 | +def _prioritize_site_packages(site_packages: Path) -> None: |
| 96 | + # Queue dependencies may be exposed through PYTHONPATH. Keep that environment |
| 97 | + # unchanged for Cinema 4D, but make the attached wheels win module resolution. |
| 98 | + (site_packages / "_deadline_adaptor_override.pth").write_text( |
| 99 | + f"import sys; sys.path.insert(0, {str(site_packages)!r})\n", |
| 100 | + encoding="utf8", |
| 101 | + ) |
| 102 | + |
| 103 | + |
| 104 | +def _emit_environment_changes( |
| 105 | + before: dict[str, str], |
| 106 | + *, |
| 107 | + venv_dir: Path, |
| 108 | + venv_bin: Path, |
| 109 | +) -> None: |
| 110 | + after = dict(before) |
| 111 | + after["PATH"] = os.pathsep.join(filter(None, (str(venv_bin), before.get("PATH", "")))) |
| 112 | + after["VIRTUAL_ENV"] = str(venv_dir) |
| 113 | + after.pop("PYTHONHOME", None) |
| 114 | + |
| 115 | + for key, value in sorted(after.items()): |
| 116 | + if value != before.get(key): |
| 117 | + print(f"openjd_env: {key}={value}") |
| 118 | + |
| 119 | + for key in sorted(before): |
| 120 | + if key not in after: |
| 121 | + print(f"openjd_unset_env: {key}") |
| 122 | + |
| 123 | + |
| 124 | +def _parse_args(argv: Sequence[str] | None = None) -> tuple[Path, Path, str]: |
| 125 | + parser = argparse.ArgumentParser() |
| 126 | + parser.add_argument("working_directory") |
| 127 | + parser.add_argument("wheels_directory") |
| 128 | + parser.add_argument("adaptor_name") |
| 129 | + args = parser.parse_args(argv) |
| 130 | + return Path(args.working_directory), Path(args.wheels_directory), args.adaptor_name |
| 131 | + |
| 132 | + |
| 133 | +def main(argv: Sequence[str] | None = None) -> None: |
| 134 | + working_dir, wheels_dir, adaptor_name = _parse_args(argv) |
| 135 | + before_environment = dict(os.environ) |
| 136 | + |
| 137 | + print(f"Setting up {adaptor_name} from attached wheels on {sys.platform}") |
| 138 | + wheels = _find_wheels(wheels_dir) |
| 139 | + for wheel in wheels: |
| 140 | + print(f" {wheel.name}") |
| 141 | + |
| 142 | + venv_dir = working_dir / "adaptor-venv" |
| 143 | + print(f"Creating adaptor virtual environment at {venv_dir}") |
| 144 | + subprocess.run( |
| 145 | + [sys.executable, "-m", "venv", "--system-site-packages", str(venv_dir)], |
| 146 | + check=True, |
| 147 | + ) |
| 148 | + |
| 149 | + venv_bin, venv_python, adaptor_executable = _get_venv_paths(venv_dir, adaptor_name) |
| 150 | + _install_wheels(venv_python, wheels) |
| 151 | + |
| 152 | + if not adaptor_executable.is_file(): |
| 153 | + raise RuntimeError( |
| 154 | + f"The override adaptor '{adaptor_name}' was not installed at {adaptor_executable}" |
| 155 | + ) |
| 156 | + |
| 157 | + site_packages = _get_site_packages(venv_python) |
| 158 | + _prioritize_site_packages(site_packages) |
| 159 | + adaptor_version = _get_distribution_version(site_packages, "deadline-cloud-for-cinema-4d") |
| 160 | + print( |
| 161 | + "ADAPTOR_OVERRIDE_READY " |
| 162 | + f"executable={adaptor_executable} " |
| 163 | + f"package=deadline-cloud-for-cinema-4d version={adaptor_version}" |
| 164 | + ) |
| 165 | + |
| 166 | + _emit_environment_changes( |
| 167 | + before_environment, |
| 168 | + venv_dir=venv_dir, |
| 169 | + venv_bin=venv_bin, |
| 170 | + ) |
| 171 | + |
| 172 | + |
| 173 | +if __name__ == "__main__": |
| 174 | + try: |
| 175 | + main() |
| 176 | + except Exception as exc: |
| 177 | + print(f"ADAPTOR_OVERRIDE_FAILED: {exc}", file=sys.stderr) |
| 178 | + raise |
0 commit comments