feat(ray-jobs): support runtime environment extensions

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

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,20 +82,27 @@ 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
             ``pyproject.toml`` automatically.
         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

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(

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()
 

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",