""" Repository rule to manage hermetic Python interpreter under Bazel. Version can be set via build parameter "--repo_env=HERMETIC_PYTHON_VERSION=3.11" To set wheel name, add "--repo_env=WHEEL_NAME=tensorflow_cpu" """ DEFAULT_VERSION = "3.11" def _python_repository_impl(ctx): version, py_kind = _get_python_version(ctx) version_and_kind = "%s-%s" % (version, py_kind) if py_kind else version ctx.file("BUILD", "") wheel_name = ctx.os.environ.get("WHEEL_NAME", "tensorflow") wheel_collab = ctx.os.environ.get("WHEEL_COLLAB", False) macos_deployment_target = ctx.os.environ.get("MACOSX_DEPLOYMENT_TARGET", "") hermetic_url = ctx.os.environ.get("HERMETIC_PYTHON_URL", "") hermetic_sha256 = ctx.os.environ.get("HERMETIC_PYTHON_SHA256", "") hermetic_prefix = ctx.os.environ.get("HERMETIC_PYTHON_PREFIX", "python") custom_requirements = ctx.os.environ.get("HERMETIC_REQUIREMENTS_LOCK", None) local_wheel_overrides_label = "" if not (hermetic_url + hermetic_sha256) and (hermetic_url or hermetic_sha256): fail(""" Please either specify both HERMETIC_PYTHON_URL and HERMETIC_PYTHON_SHA256 to set up a custom python interpreter, or none of them to rely on default ones. """) requirements = None if not requirements: for i in range(0, len(ctx.attr.requirements_locks)): if ctx.attr.requirements_versions[i] == version_and_kind: requirements = ctx.attr.requirements_locks[i] break if not requirements and not custom_requirements: fail(""" Could not find requirements_lock.txt file matching specified Python version. Specified python version: {version} Python versions with available requirement_lock.txt files: {versions} Please check python_init_repositories() in your WORKSPACE file. """.format( version = version, versions = ", ".join(ctx.attr.requirements_versions), )) if custom_requirements: custom_requirements_path = ctx.path(custom_requirements) requirements_with_local_wheels = "@{repo}//:{label}".format( repo = ctx.name, label = custom_requirements_path.basename, ) ctx.file( custom_requirements_path.basename, ctx.read(custom_requirements_path), ) elif ctx.attr.local_wheel_workspaces: base_requirements = ctx.read(requirements) merged_requirements_blocks, override_report_entries = ( _rewrite_requirements_with_local_wheels( ctx, version, ctx.attr.local_wheel_workspaces, base_requirements, ) ) merged_requirements_content = "\n".join(merged_requirements_blocks) requirements_with_local_wheels = "@{repo}//:{label}".format( repo = ctx.name, label = requirements.name, ) ctx.file( requirements.name, merged_requirements_content, ) local_wheel_overrides_label = _write_local_wheel_overrides_json( ctx, version, py_kind, requirements_with_local_wheels, override_report_entries, ) _print_local_wheel_override_summaries(override_report_entries) else: requirements_with_local_wheels = str(requirements) use_pywrap_rules = bool( ctx.os.environ.get("USE_PYWRAP_RULES", False), ) if use_pywrap_rules: print("!!!Using pywrap rules instead of directly creating .so objects!!!") # buildifier: disable=print interpreter_type = "\"default\" (provided by rules_python)" if hermetic_url: interpreter_type = "\"custom\" (pulled from %s)" % hermetic_url print( """ ============================= Hermetic Python configuration: Version: "{version}" Kind: "{py_kind}" Interpreter: {interpreter_type} Requirements_lock label: "{requirements_lock_label}" ===================================== """.format( version = version, py_kind = py_kind, interpreter_type = interpreter_type, requirements_lock_label = requirements_with_local_wheels, ), ) # buildifier: disable=print ctx.file( "py_version.bzl", """ TF_PYTHON_VERSION = "{version}" HERMETIC_PYTHON_VERSION = "{version}" HERMETIC_PYTHON_VERSION_KIND = "{py_kind}" WHEEL_NAME = "{wheel_name}" WHEEL_COLLAB = "{wheel_collab}" REQUIREMENTS = "{requirements}" REQUIREMENTS_WITH_LOCAL_WHEELS = "{requirements_with_local_wheels}" LOCAL_WHEEL_OVERRIDES_LABEL = "{local_wheel_overrides_label}" USE_PYWRAP_RULES = {use_pywrap_rules} MACOSX_DEPLOYMENT_TARGET = "{macos_deployment_target}" HERMETIC_PYTHON_URL = "{hermetic_url}" HERMETIC_PYTHON_SHA256 = "{hermetic_sha256}" HERMETIC_PYTHON_PREFIX = "{hermetic_prefix}" """.format( version = version, py_kind = py_kind, wheel_name = wheel_name, wheel_collab = wheel_collab, requirements = str(requirements), requirements_with_local_wheels = requirements_with_local_wheels, local_wheel_overrides_label = local_wheel_overrides_label, use_pywrap_rules = use_pywrap_rules, macos_deployment_target = macos_deployment_target, hermetic_url = hermetic_url, hermetic_sha256 = hermetic_sha256, hermetic_prefix = hermetic_prefix, ), ) def _get_python_version(ctx): print_warning = False version = ctx.os.environ.get("HERMETIC_PYTHON_VERSION", "") if not version: version = ctx.os.environ.get("TF_PYTHON_VERSION", "") if not version: print_warning = True if ctx.attr.default_python_version == "system": python_version_result = ctx.execute(["python3", "--version"]) if python_version_result.return_code == 0: version = python_version_result.stdout else: fail(""" Cannot match hermetic Python version to system Python version. System Python was not found.""") else: version = ctx.attr.default_python_version version, kind = _parse_python_version(version) if print_warning: print(""" HERMETIC_PYTHON_VERSION variable was not set correctly, using default version. Python {} will be used. To select Python version, either set HERMETIC_PYTHON_VERSION env variable in your shell: export HERMETIC_PYTHON_VERSION=3.12 OR pass it as an argument to bazel command directly or inside your .bazelrc file: --repo_env=HERMETIC_PYTHON_VERSION=3.12 """.format(version)) # buildifier: disable=print return version, kind def _parse_python_version(version_str): if version_str.startswith("Python "): py_ver_chunks = version_str[7:].split(".") return "%s.%s" % (py_ver_chunks[0], py_ver_chunks[1]), "" elif "-" in version_str: return version_str.split("-") return version_str, "" def _rewrite_requirements_with_local_wheels( ctx, py_version, local_wheel_workspaces, base_requirements): """Rewrites the selected lockfile so matching packages resolve from local wheels.""" os_name = ctx.os.name is_windows = "windows" in os_name.lower() local_file_path_prefix = "file:" if is_windows else "file://" # Parse the lockfile once so later helpers can work with package-level data # instead of re-scanning the raw text in multiple places. base_requirement_blocks = _split_requirements_blocks(base_requirements) base_requirement_names = _collect_requirement_names(base_requirement_blocks) local_wheels = _collect_local_wheels( ctx, py_version, local_wheel_workspaces, base_requirement_names, ) merged_requirements_blocks, override_report_entries = _merge_local_wheels_into_base_requirements( base_requirement_blocks, local_wheels, local_file_path_prefix, ) return merged_requirements_blocks, override_report_entries def _collect_local_wheels( ctx, py_version, local_wheel_workspaces, base_requirement_names): """Collects the newest matching local wheel for each distribution name.""" py_ver_marker = "-cp%s-" % py_version.replace(".", "") py_major_ver_marker = "-py%s-" % py_version.split(".")[0] wheels = {} if local_wheel_workspaces: for local_wheel_workspace in local_wheel_workspaces: local_wheel_workspace_path = ctx.path(local_wheel_workspace) dist_folder = ctx.attr.local_wheel_dist_folder dist_folder_path = local_wheel_workspace_path.dirname.get_child(dist_folder) if dist_folder_path.exists: dist_wheels = dist_folder_path.readdir() _process_dist_wheels( dist_wheels, wheels, py_ver_marker, py_major_ver_marker, ctx.attr.local_wheel_inclusion_list, ctx.attr.local_wheel_exclusion_list, ) local_wheels = {} for wheel_name, wheel_path in wheels.items(): # Normalize `foo_bar` to `foo-bar`. We assume that, if `foo_bar` # isn't present in requirements, it must be named `foo-bar`. The # exact same distribution name needs to be used to ensure it is # correctly overridden. if "_" in wheel_name and wheel_name not in base_requirement_names: local_package_name = wheel_name.replace("_", "-") else: local_package_name = wheel_name local_wheels[_normalize_requirement_name(local_package_name)] = { # Keep the concrete wheel path so the merge step can emit an # explicit `pkg @ file://...whl` direct reference. "wheel_path": wheel_path, } return local_wheels def _merge_local_wheels_into_base_requirements( base_requirement_blocks, local_wheels, local_file_path_prefix): """Rewrites matching requirement blocks to use direct local wheel URLs.""" # Replace matching lockfile entries with local wheel entries, so they take # precedence in both host-platform and download_only resolution paths. merged_blocks = [] override_report_entries = [] replaced_packages = {} for block in base_requirement_blocks: first_line = block[0] if block else "" package_name = _extract_requirement_name(first_line) if package_name and package_name in local_wheels: local_wheel = local_wheels[package_name] marker = _extract_requirement_marker(first_line) comment_lines = [] for line in block[1:]: # Keep explanatory comments such as "# via ...", but drop the # old PyPI hash continuations because the replacement now points # pip at a concrete local wheel URL. if line.strip().startswith("#"): comment_lines.append(line) rendered_block, override_report_entry = _make_local_wheel_override( package_name = package_name, local_wheel = local_wheel, marker = marker, comment_lines = comment_lines, local_file_path_prefix = local_file_path_prefix, ) merged_blocks.append(rendered_block) override_report_entries.append(override_report_entry) replaced_packages[package_name] = True else: merged_blocks.append("\n".join(block)) for package_name, local_wheel in local_wheels.items(): if package_name not in replaced_packages: # Keep the old behavior of surfacing local wheels even when the base # lockfile has no entry for the package yet. rendered_block, override_report_entry = _make_local_wheel_override( package_name = package_name, local_wheel = local_wheel, marker = "", comment_lines = [], local_file_path_prefix = local_file_path_prefix, ) merged_blocks.append(rendered_block) override_report_entries.append(override_report_entry) return merged_blocks, override_report_entries def _print_local_wheel_override_summaries(override_report_entries): if not override_report_entries: return print( """ Local wheel overrides: {overrides} """.format( overrides = "\n".join([ _format_local_wheel_override_summary(override_report_entry) for override_report_entry in override_report_entries ]), ), ) # buildifier: disable=print def _write_local_wheel_overrides_json( ctx, version, py_kind, requirements_with_local_wheels, override_report_entries): label = "@{repo}//:local_wheel_overrides.json".format(repo = ctx.name) ctx.file( "local_wheel_overrides.json", json.encode_indent({ "python_version": version, "python_kind": py_kind, "requirements_lock_label": requirements_with_local_wheels, "overrides": override_report_entries, }), ) return label def _split_requirements_blocks(requirements_content): """Groups a pip-compile-style requirements file into top-level entries.""" blocks = [] current_block = [] for line in requirements_content.replace("\r", "").split("\n"): if current_block and not _is_continuation_line(line): blocks.append(current_block) current_block = [] current_block.append(line) if current_block: blocks.append(current_block) return blocks def _is_continuation_line(line): return line.startswith(" ") or line.startswith("\t") def _collect_requirement_names(requirement_blocks): """Collects normalized requirement names from lockfile head lines.""" requirement_names = {} for block in requirement_blocks: first_line = block[0] if block else "" requirement_name = _extract_requirement_name(first_line) if requirement_name: requirement_names[requirement_name] = True return requirement_names def _extract_requirement_name(line): """Extracts the normalized package name from a requirement head line.""" stripped_line = line.strip() if not stripped_line or stripped_line.startswith("#") or stripped_line.startswith("-"): return None package_name = "" for character in stripped_line.elems(): if character.isspace() or character in [">", "<", "~", "=", ";", "[", "@"]: break package_name += character if not package_name: return None return _normalize_requirement_name(package_name) def _normalize_requirement_name(package_name): return package_name.replace("_", "-").lower() def _extract_requirement_marker(line): """Preserves the raw "; ..." condition that controls when a requirement applies. For example, this returns 'sys_platform == "linux"' for 'jax-cuda12-plugin==0.9.2 ; sys_platform == "linux"'. """ stripped_line = line.strip() if stripped_line.endswith("\\"): stripped_line = stripped_line[:-1].rstrip() # If the line doesn't contain the separator, partition(";") returns (line, "", ""). _, marker_separator, marker = stripped_line.partition(";") if not marker_separator: return "" return marker.strip() def _make_local_wheel_override( package_name, local_wheel, marker, comment_lines, local_file_path_prefix): """Builds both the rewritten requirement block and its reporting metadata.""" rendered_block = _render_local_override_requirement_block( package_name = package_name, wheel_path = local_wheel["wheel_path"], marker = marker, comment_lines = comment_lines, local_file_path_prefix = local_file_path_prefix, ) override_report_entry = _make_local_wheel_override_entry( package_name = package_name, wheel_path = local_wheel["wheel_path"], marker = marker, ) return rendered_block, override_report_entry def _render_local_override_requirement_block( package_name, wheel_path, marker, comment_lines, local_file_path_prefix): requirement_line = "{package_name} @ {local_file_path_prefix}{wheel_path}".format( package_name = package_name, local_file_path_prefix = local_file_path_prefix, wheel_path = wheel_path.realpath, ) if marker: requirement_line += " ; " + marker rendered_lines = [requirement_line] rendered_lines.extend(comment_lines) return "\n".join(rendered_lines) def _make_local_wheel_override_entry(package_name, wheel_path, marker): return { "package": package_name, "marker": marker, "wheel": wheel_path.basename, } def _format_local_wheel_override_summary(override_entry): summary = " {package_name}".format(package_name = override_entry["package"]) if override_entry["marker"]: summary += " ; " + override_entry["marker"] return "{summary} -> {wheel_name}".format( summary = summary, wheel_name = override_entry["wheel"], ) python_repository = repository_rule( implementation = _python_repository_impl, attrs = { "requirements_versions": attr.string_list( mandatory = False, default = [], ), "requirements_locks": attr.label_list( mandatory = False, default = [], ), "local_wheel_workspaces": attr.label_list( mandatory = False, default = [], ), "local_wheel_dist_folder": attr.string( mandatory = False, default = "dist", ), "default_python_version": attr.string( mandatory = False, default = DEFAULT_VERSION, ), "local_wheel_inclusion_list": attr.string_list( mandatory = False, default = ["*"], ), "local_wheel_exclusion_list": attr.string_list( mandatory = False, default = [], ), }, environ = [ "TF_PYTHON_VERSION", "HERMETIC_PYTHON_VERSION", "HERMETIC_PYTHON_URL", "HERMETIC_PYTHON_SHA256", "HERMETIC_REQUIREMENTS_LOCK", "HERMETIC_PYTHON_PREFIX", "WHEEL_NAME", "WHEEL_COLLAB", "USE_PYWRAP_RULES", "MACOSX_DEPLOYMENT_TARGET", ], local = True, ) def _process_dist_wheels( dist_wheels, wheels, py_ver_marker, py_major_ver_marker, local_wheel_inclusion_list, local_wheel_exclusion_list): for wheel in dist_wheels: bn = wheel.basename if not bn.endswith(".whl") or (bn.find(py_ver_marker) < 0 and bn.find(py_major_ver_marker) < 0): continue if not _basic_wildcard_match(bn, local_wheel_inclusion_list, True, False): continue if not _basic_wildcard_match(bn, local_wheel_exclusion_list, False, True): continue name_components = bn.split("-") package_name = name_components[0] for name_component in name_components[1:]: if name_component[0].isdigit(): break package_name += "-" + name_component latest_wheel = wheels.get(package_name, None) if not latest_wheel or latest_wheel.basename < wheel.basename: wheels[package_name] = wheel def _basic_wildcard_match(name, patterns, expected_match_result, match_all): match = False for pattern in patterns: match = False if pattern.startswith("*") and pattern.endswith("*"): match = name.find(pattern[1:-1]) >= 0 elif pattern.startswith("*"): match = name.endswith(pattern[1:]) elif pattern.endswith("*"): match = name.startswith(pattern[:-1]) else: match = name == pattern if match_all: if match != expected_match_result: return False elif match == expected_match_result: return True return match == expected_match_result def _custom_python_interpreter_impl(ctx): version = ctx.attr.version version_variant = ctx.attr.version_variant strip_prefix = ctx.attr.strip_prefix.format( version = version, version_variant = version_variant, ) urls = [url.format(version = version, version_variant = version_variant) for url in ctx.attr.urls] binary_name = ctx.attr.binary_name if not binary_name: ver_chunks = version.split(".") binary_name = "python%s.%s" % (ver_chunks[0], ver_chunks[1]) install_dir = "{name}-{version}".format(name = ctx.attr.name, version = version) _exec_and_check(ctx, ["mkdir", install_dir]) install_path = ctx.path(install_dir) srcs_dir = "srcs" ctx.download_and_extract( url = urls, stripPrefix = strip_prefix, output = srcs_dir, ) configure_params = list(ctx.attr.configure_params) if "CC" in ctx.os.environ: configure_params.append("CC={}".format(ctx.os.environ["CC"])) if "CXX" in ctx.os.environ: configure_params.append("CXX={}".format(ctx.os.environ["CXX"])) configure_params.append("--prefix=%s" % install_path.realpath) _exec_and_check( ctx, ["./configure"] + configure_params, working_directory = srcs_dir, quiet = False, ) res = _exec_and_check(ctx, ["nproc"]) cores = 12 if res.return_code != 0 else max(1, int(res.stdout.strip()) - 1) _exec_and_check(ctx, ["make", "-j%s" % cores], working_directory = srcs_dir) _exec_and_check(ctx, ["make", "altinstall"], working_directory = srcs_dir) _exec_and_check(ctx, ["ln", "-s", binary_name, "python3"], working_directory = install_dir + "/bin") tar = "{install_dir}.tgz".format(install_dir = install_dir) _exec_and_check(ctx, ["tar", "czpf", tar, install_dir]) _exec_and_check(ctx, ["rm", "-rf", srcs_dir]) res = _exec_and_check(ctx, ["sha256sum", tar]) sha256 = res.stdout.split(" ")[0].strip() tar_path = ctx.path(tar) example = """\n\n To use newly built Python interpreter add the following code snippet RIGHT AFTER python_init_toolchains() in your WORKSPACE file. The code sample should work as is but it may need some tuning, if you have special requirements. ``` load("@rules_python//python:repositories.bzl", "python_register_toolchains") python_register_toolchains( name = "python", # By default assume the interpreter is on the local file system, replace # with proper URL if it is not the case. base_url = "file://", ignore_root_user_error = True, python_version = "{version}", tool_versions = {{ "{version}": {{ # Path to .tar.gz with Python binary. By default it points to .tgz # file in cache where it was built originally; replace with proper # file location, if you moved it somewhere else. "url": "{tar_path}", "sha256": {{ # By default we assume Linux x86_64 architecture, eplace with # proper architecture if you were building on a different platform. "x86_64-unknown-linux-gnu": "{sha256}", }}, "strip_prefix": "{install_dir}", }}, }}, ) ``` \n\n""".format(version = version, tar_path = tar_path, sha256 = sha256, install_dir = install_dir) instructions = "INSTRUCTIONS-{version}.md".format(version = version) ctx.file(instructions + ".tmpl", example, executable = False) ctx.file( "BUILD.bazel", """ genrule( name = "{name}", srcs = ["{tar}", "{instructions}.tmpl"], outs = ["{install_dir}.tar.gz", "{instructions}"], cmd = "cp $(location {tar}) $(location {install_dir}.tar.gz); cp $(location {instructions}.tmpl) $(location {instructions})", visibility = ["//visibility:public"], ) """.format( name = ctx.attr.name, tar = tar, install_dir = install_dir, instructions = instructions, ), executable = False, ) print(example) # buildifier: disable=print custom_python_interpreter = repository_rule( implementation = _custom_python_interpreter_impl, attrs = { "urls": attr.string_list(), "strip_prefix": attr.string(), "binary_name": attr.string(mandatory = False), "version": attr.string(), "version_variant": attr.string(), "configure_params": attr.string_list( mandatory = False, default = ["--enable-optimizations"], ), }, ) def _exec_and_check(ctx, command, fail_on_error = True, quiet = False, **kwargs): res = ctx.execute(command, quiet = quiet, **kwargs) if fail_on_error and res.return_code != 0: fail(""" Failed to execute command: `{command}` Exit Code: {code} STDERR: {stderr} """.format( command = command, code = res.return_code, stderr = res.stderr, )) return res