diff --git a/docs/adr/002-broker-plugin-boundary.md b/docs/adr/002-broker-plugin-boundary.md index f338d33..5cdd30e 100644 --- a/docs/adr/002-broker-plugin-boundary.md +++ b/docs/adr/002-broker-plugin-boundary.md @@ -50,6 +50,13 @@ that driver calls existing in-process training or batch-inference APIs. Worker side broker cancellation, arbitrary execution context, a generic Core consume loop, and durable workflow semantics are outside the Alpha contract. +Trusted deployment configuration may supply an execution-driver package through +the generic Ray Jobs runtime-environment arguments on `submit_ray_job`. Core +only forwards explicitly declared modules and requirements; it does not inspect +broker payloads, discover provider installations, or resolve dependencies. This +deployment mechanism is separate from Broker API v1 and does not change +`BROKER_API_VERSION`. + `submission_id` is the primary Ray Jobs identity for admission, status, logs, and stop. `ray_job_id` is optional execution metadata and is populated only from a real Ray `JobDetails.job_id`; Core never substitutes `submission_id` for diff --git a/src/tributo/_common/runtime_env.py b/src/tributo/_common/runtime_env.py index 1d633b1..696886e 100644 --- a/src/tributo/_common/runtime_env.py +++ b/src/tributo/_common/runtime_env.py @@ -1,8 +1,9 @@ """Ray runtime environment builder. Ships application code through Ray-managed ``working_dir`` and ``py_modules``. -Python package dependencies remain owned by the cluster image or an explicit -Ray runtime environment instead of being merged from unrelated environments. +Python package dependencies remain owned by the cluster image, trusted runtime +configuration, or an explicit algorithm artifact instead of being merged from +unrelated environments. """ from __future__ import annotations @@ -66,6 +67,8 @@ def build_runtime_env( project_root: Path | None = None, package_name: str = "tributo", extra_excludes: list[str] | None = None, + extra_py_modules: list[str | Path] | None = None, + runtime_pip_packages: list[str] | None = None, env_vars: dict[str, str] | None = None, pythonpath: str | None = None, algorithm_artifact: AlgorithmArtifact | None = None, @@ -79,13 +82,14 @@ def build_runtime_env( older version baked into the image; 2. ``working_dir`` uploads the project root, providing entrypoint scripts; - Package dependencies must be installed in the cluster image or supplied by - an explicit, preflighted offline Ray runtime environment. When - ``algorithm_artifact`` is present, the builder validates the Wheel or - Bundle against ``image_profile`` and owns the artifact-related - ``working_dir``, ``py_modules``, ``pip``, and ``excludes`` fields. It never - derives package paths from the submitting process or combines unrelated - ``site-packages`` trees. + Package dependencies must be installed in the cluster image, supplied by + trusted deployment configuration, or distributed through an explicit, + preflighted algorithm artifact. Extension inputs must not be derived from + an untrusted task payload. Core copies them without scanning the submitting + process or resolving dependencies. When ``algorithm_artifact`` is present, + the builder validates the Wheel or Bundle against ``image_profile`` and + owns the artifact-related ``working_dir``, ``py_modules``, ``pip``, and + ``excludes`` fields. Args: project_root: Project root directory. When ``None``, walks up to find @@ -93,6 +97,12 @@ def build_runtime_env( package_name: Package name to upload for priority override, defaults to ``tributo``. extra_excludes: Additional directories/files to exclude (glob patterns). + extra_py_modules: Trusted extension modules appended after the Tributo + Core package in caller-provided order. Paths are converted to + strings. Do not populate this from a broker task payload. + runtime_pip_packages: Trusted extension requirements copied to Ray's + ``pip`` runtime environment. Core does not resolve them. This must + not be combined with ``algorithm_artifact``. env_vars: Additional environment variables. pythonpath: Optional cluster-visible ``PYTHONPATH`` to append to an explicitly supplied value. ``None`` leaves ``PYTHONPATH`` untouched. @@ -109,6 +119,10 @@ def build_runtime_env( A dict that can be passed directly to ``JobSubmissionClient.submit_job(runtime_env=...)``. + Raises: + ValueError: ``runtime_pip_packages`` and ``algorithm_artifact`` both + request ownership of Ray's ``pip`` runtime environment. + Example: >>> from tributo._common.runtime_env import build_runtime_env >>> runtime_env = build_runtime_env() @@ -117,6 +131,11 @@ def build_runtime_env( ... runtime_env=runtime_env, ... ) """ + if runtime_pip_packages and algorithm_artifact is not None: + raise ValueError( + "runtime_pip_packages cannot be combined with algorithm_artifact" + ) + if project_root is None: try: root = find_project_root() @@ -158,10 +177,15 @@ def build_runtime_env( runtime_env: dict[str, Any] = { "working_dir": str(root), "excludes": excludes, - "py_modules": [str(src_pkg)], + "py_modules": [ + str(src_pkg), + *(str(module) for module in extra_py_modules or ()), + ], } if merged_env_vars: runtime_env["env_vars"] = merged_env_vars + if runtime_pip_packages: + runtime_env["pip"] = list(runtime_pip_packages) if algorithm_artifact is None and ( image_profile is not None or declared_dependencies @@ -204,9 +228,9 @@ def build_runtime_env( runtime_env["env_vars"] = merged_env_vars logger.debug( - "Built runtime_env: working_dir=%s, py_modules=%s, env_var_keys=%s", - runtime_env["working_dir"], - runtime_env["py_modules"], + "Built runtime_env: py_module_count=%d, has_pip=%s, env_var_keys=%s", + len(runtime_env["py_modules"]), + "pip" in runtime_env, sorted(merged_env_vars), ) return runtime_env diff --git a/src/tributo/ray_jobs.py b/src/tributo/ray_jobs.py index d61d82a..18b8dec 100644 --- a/src/tributo/ray_jobs.py +++ b/src/tributo/ray_jobs.py @@ -157,13 +157,21 @@ def submit_ray_job( env_vars: dict[str, str] | None = None, project_root: Path | None = None, extra_excludes: list[str] | None = None, + extra_py_modules: list[str | Path] | None = None, + runtime_pip_packages: list[str] | None = None, metadata: dict[str, str] | None = None, request_digest: str | None = None, entrypoint_num_cpus: float | None = None, entrypoint_num_gpus: float | None = None, entrypoint_memory: int | None = None, ) -> RayJobSubmission: - """Submit one deterministic Ray Job and reconcile an ambiguous response.""" + """Submit one deterministic Ray Job and reconcile an ambiguous response. + + ``extra_py_modules`` and ``runtime_pip_packages`` are trusted deployment + configuration, not broker task payload fields. Core forwards these values + without discovering providers or resolving dependencies. Runtime pip + packages cannot be combined with algorithm artifact pip distribution. + """ if not entrypoint.strip(): raise ValueError("entrypoint must not be empty") @@ -192,6 +200,8 @@ def submit_ray_job( project_root=project_root, env_vars=job_env, extra_excludes=extra_excludes, + extra_py_modules=extra_py_modules, + runtime_pip_packages=runtime_pip_packages, ) client = _get_submission_client(dashboard_url) return _submit_ray_job_with_client( diff --git a/tests/algorithms/test_algorithm_distribution.py b/tests/algorithms/test_algorithm_distribution.py index 6f8da6e..bfc0ee0 100644 --- a/tests/algorithms/test_algorithm_distribution.py +++ b/tests/algorithms/test_algorithm_distribution.py @@ -379,17 +379,22 @@ def test_offline_wheelhouse_requires_image_pip(tmp_path: Path) -> None: prepare_algorithm_distribution(artifact, profile) -def test_build_runtime_env_combines_tributo_and_code_wheel(tmp_path: Path) -> None: +def test_build_runtime_env_orders_core_extension_and_code_wheel( + tmp_path: Path, +) -> None: wheel = _write_wheel(tmp_path, "demo-algorithm") artifact = AlgorithmArtifact(source=str(wheel)) + extension_module = tmp_path / "execution-driver.whl" runtime_env = build_runtime_env( algorithm_artifact=artifact, image_profile=_profile(), + extra_py_modules=[extension_module], ) - assert len(runtime_env["py_modules"]) == 2 - assert runtime_env["py_modules"][-1] == str(wheel) + assert len(runtime_env["py_modules"]) == 3 + assert Path(runtime_env["py_modules"][0]).name == "tributo" + assert runtime_env["py_modules"][-2:] == [str(extension_module), str(wheel)] working_dir = Path(runtime_env["working_dir"]) assert (working_dir / "pyproject.toml").is_file() diff --git a/tests/test_ray_jobs.py b/tests/test_ray_jobs.py index 7c180bb..03245fd 100644 --- a/tests/test_ray_jobs.py +++ b/tests/test_ray_jobs.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import Any +from pathlib import Path from unittest.mock import MagicMock, patch import pytest @@ -17,11 +17,6 @@ ) -def _runtime_env(*args: Any, **kwargs: Any) -> dict[str, Any]: - del args - return {"env_vars": kwargs.get("env_vars", {})} - - def test_submission_identity_is_workload_neutral_and_ray_job_id_is_real() -> None: client = MagicMock() client.submit_job.return_value = "ray-api-return-value" @@ -31,7 +26,11 @@ def test_submission_identity_is_workload_neutral_and_ray_job_id_is_real() -> Non with ( patch("tributo.ray_jobs._get_submission_client", return_value=client), - patch("tributo.ray_jobs.build_runtime_env", side_effect=_runtime_env), + patch( + "tributo.ray_jobs.build_runtime_env", + autospec=True, + return_value={"env_vars": {}}, + ), ): result = submit_ray_job( "python -m provider.driver", @@ -48,6 +47,73 @@ def test_submission_identity_is_workload_neutral_and_ray_job_id_is_real() -> Non assert client.submit_job.call_args.kwargs["submission_id"] == result.submission_id +def test_submit_ray_job_forwards_trusted_runtime_env_extensions( + tmp_path: Path, +) -> None: + client = MagicMock() + client.get_job_info.return_value = type( + "JobInfo", (), {"job_id": "ray-core-job-2"} + )() + extensions: list[str | Path] = [ + tmp_path / "driver.whl", + "s3://artifacts/support.zip", + ] + packages = ["driver-runtime==1.2.3", "/artifacts/support.whl"] + built_runtime_env = { + "py_modules": [str(tmp_path / "tributo"), *map(str, extensions)], + "pip": list(packages), + } + + with ( + patch("tributo.ray_jobs._get_submission_client", return_value=client), + patch( + "tributo.ray_jobs.build_runtime_env", + autospec=True, + return_value=built_runtime_env, + ) as build_runtime_env_mock, + ): + result = submit_ray_job( + "python -m extension.execution_driver", + operation_namespace="broker", + run_id="run-2", + attempt_id="attempt-3", + env_vars={"DEPLOYMENT_MODE": "trusted"}, + project_root=tmp_path, + extra_excludes=["local-cache/**"], + extra_py_modules=extensions, + runtime_pip_packages=packages, + metadata={"operation_type": "training"}, + request_digest="digest-2", + entrypoint_num_cpus=1.5, + entrypoint_num_gpus=0.5, + entrypoint_memory=1024, + ) + + build_runtime_env_mock.assert_called_once_with( + project_root=tmp_path, + env_vars={ + "DEPLOYMENT_MODE": "trusted", + "TRIBUTO_RUN_ID": "run-2", + "TRIBUTO_ATTEMPT_ID": "attempt-3", + "TRIBUTO_SUBMISSION_ID": result.submission_id, + }, + extra_excludes=["local-cache/**"], + extra_py_modules=extensions, + runtime_pip_packages=packages, + ) + submit_call = client.submit_job.call_args.kwargs + assert submit_call["runtime_env"] is built_runtime_env + assert submit_call["metadata"] == { + "operation_type": "training", + "tributo.request_digest": "digest-2", + } + assert submit_call["submission_id"] == result.submission_id + assert submit_call["entrypoint_num_cpus"] == 1.5 + assert submit_call["entrypoint_num_gpus"] == 0.5 + assert submit_call["entrypoint_memory"] == 1024 + assert result.ray_job_id == "ray-core-job-2" + + def test_ambiguous_submission_reconciles_by_submission_id() -> None: client = MagicMock() client.submit_job.side_effect = TimeoutError("response lost") @@ -56,7 +122,11 @@ def test_ambiguous_submission_reconciles_by_submission_id() -> None: with ( patch("tributo.ray_jobs._get_submission_client", return_value=client), - patch("tributo.ray_jobs.build_runtime_env", side_effect=_runtime_env), + patch( + "tributo.ray_jobs.build_runtime_env", + autospec=True, + return_value={"env_vars": {}}, + ), ): result = submit_ray_job( "python -m provider.driver", @@ -74,7 +144,11 @@ def test_request_digest_is_optional_metadata_not_submission_identity() -> None: with ( patch("tributo.ray_jobs._get_submission_client", return_value=client), - patch("tributo.ray_jobs.build_runtime_env", side_effect=_runtime_env), + patch( + "tributo.ray_jobs.build_runtime_env", + autospec=True, + return_value={"env_vars": {}}, + ), ): first = submit_ray_job( "python -m provider.driver", diff --git a/tests/test_runtime_env.py b/tests/test_runtime_env.py index 61426be..271c01f 100644 --- a/tests/test_runtime_env.py +++ b/tests/test_runtime_env.py @@ -3,10 +3,12 @@ from __future__ import annotations import logging +from pathlib import Path import pytest -from tributo._common.runtime_env import build_runtime_env +from tributo._common.runtime_env import DEFAULT_EXCLUDES, build_runtime_env +from tributo.algorithms.api.artifacts import AlgorithmArtifact def test_runtime_env_does_not_inject_python_package_paths_by_default() -> None: @@ -38,15 +40,23 @@ def test_runtime_env_debug_log_never_exposes_environment_values( caplog: pytest.LogCaptureFixture, ) -> None: profile_payload = "opaque-profile-payload" + extension_module = "/deployment/private-extension-module" + runtime_package = "private-extension-package==9.9.9" with caplog.at_level(logging.DEBUG, logger="tributo._common.runtime_env"): runtime_env = build_runtime_env( - env_vars={"TRIBUTO_STORAGE_PROFILE_MODEL": profile_payload} + env_vars={"TRIBUTO_STORAGE_PROFILE_MODEL": profile_payload}, + extra_py_modules=[extension_module], + runtime_pip_packages=[runtime_package], ) assert runtime_env["env_vars"]["TRIBUTO_STORAGE_PROFILE_MODEL"] == profile_payload + assert runtime_env["py_modules"][-1] == extension_module + assert runtime_env["pip"] == [runtime_package] assert "PYTHONPATH" not in runtime_env["env_vars"] assert profile_payload not in caplog.text + assert extension_module not in caplog.text + assert runtime_package not in caplog.text assert "TRIBUTO_STORAGE_PROFILE_MODEL" in caplog.text @@ -56,5 +66,81 @@ def test_default_runtime_env_does_not_add_extension_dependencies(tmp_path) -> No ) (tmp_path / "tributo").mkdir() runtime_env = build_runtime_env(project_root=tmp_path) - assert runtime_env["py_modules"] == [str(tmp_path / "tributo")] + assert runtime_env == { + "working_dir": str(tmp_path), + "excludes": DEFAULT_EXCLUDES, + "py_modules": [str(tmp_path / "tributo")], + } + + +def test_runtime_env_appends_explicit_extension_modules_in_order(tmp_path) -> None: + (tmp_path / "tributo").mkdir() + first_extension = tmp_path / "extensions" / "driver.whl" + second_extension = "s3://deployment-artifacts/shared-runtime.zip" + + runtime_env = build_runtime_env( + project_root=tmp_path, + extra_py_modules=[first_extension, second_extension], + ) + + assert runtime_env["py_modules"] == [ + str(tmp_path / "tributo"), + str(first_extension), + second_extension, + ] + + +def test_runtime_env_adds_only_explicit_nonempty_pip_packages(tmp_path) -> None: + (tmp_path / "tributo").mkdir() + packages = ["driver-runtime==1.2.3", "/artifacts/support.whl"] + + runtime_env = build_runtime_env( + project_root=tmp_path, + runtime_pip_packages=packages, + ) + + assert runtime_env["pip"] == packages + assert runtime_env["pip"] is not packages + + +@pytest.mark.parametrize("runtime_pip_packages", [None, []]) +def test_runtime_env_omits_empty_pip_configuration( + tmp_path: Path, + runtime_pip_packages: list[str] | None, +) -> None: + (tmp_path / "tributo").mkdir() + + runtime_env = build_runtime_env( + project_root=tmp_path, + runtime_pip_packages=runtime_pip_packages, + ) + + assert "pip" not in runtime_env + + +def test_runtime_env_rejects_competing_pip_owners(tmp_path) -> None: + artifact = AlgorithmArtifact(source=str(tmp_path / "algorithm.whl")) + + with pytest.raises( + ValueError, + match="runtime_pip_packages cannot be combined with algorithm_artifact", + ): + build_runtime_env( + project_root=tmp_path, + runtime_pip_packages=["driver-runtime==1.2.3"], + algorithm_artifact=artifact, + ) + + +def test_runtime_env_does_not_discover_host_python_paths( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + (tmp_path / "tributo").mkdir() + host_site_packages = tmp_path / "host" / "site-packages" + monkeypatch.syspath_prepend(str(host_site_packages)) + + runtime_env = build_runtime_env(project_root=tmp_path) + + assert str(host_site_packages) not in runtime_env["py_modules"] assert "pip" not in runtime_env