diff --git a/.github/workflows/ci_python.yml b/.github/workflows/ci_python.yml index 530e3d4e6..e94948b13 100644 --- a/.github/workflows/ci_python.yml +++ b/.github/workflows/ci_python.yml @@ -153,6 +153,38 @@ jobs: - name: Build wheels run: just wheels + - name: Test installed adapter discovery + run: | + set -euo pipefail + uv venv .wheel-test-venv + if [[ -x .wheel-test-venv/bin/python ]]; then + test_python="$PWD/.wheel-test-venv/bin/python" + else + test_python="$PWD/.wheel-test-venv/Scripts/python.exe" + fi + uv pip install \ + --python "$test_python" \ + --find-links dist \ + "nemo-fabric[hermes]" + + source_adapters="$GITHUB_WORKSPACE/adapters" + staged_source_adapters="$RUNNER_TEMP/nemo-fabric-source-adapters" + mv "$source_adapters" "$staged_source_adapters" + trap 'mv "$staged_source_adapters" "$source_adapters"' EXIT + + smoke_dir="$(mktemp -d)" + cd "$smoke_dir" + "$test_python" - <<'PY' + from nemo_fabric import Fabric, FabricConfig, HarnessConfig, MetadataConfig + + config = FabricConfig( + metadata=MetadataConfig(name="installed-adapter-smoke"), + harness=HarnessConfig(adapter_id="nvidia.fabric.hermes"), + ) + plan = Fabric().plan(config) + assert plan.adapter.adapter_id == "nvidia.fabric.hermes" + PY + - name: Upload wheels uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: diff --git a/crates/fabric-core/src/config.rs b/crates/fabric-core/src/config.rs index 5316c4337..b405676c2 100644 --- a/crates/fabric-core/src/config.rs +++ b/crates/fabric-core/src/config.rs @@ -164,10 +164,17 @@ struct AdapterRegistry { } impl AdapterRegistry { - fn from_config(_config: &FabricConfig, base_dir: &Path) -> Result { + fn from_config( + _config: &FabricConfig, + base_dir: &Path, + adapter_descriptors: &[PathBuf], + ) -> Result { let mut registry = Self::default(); registry.register_repository_directory(&repository_adapter_dir())?; registry.register_local_directory(&base_dir.join("adapters"))?; + for path in adapter_descriptors { + registry.register_descriptor_file(path, AdapterDescriptorSource::Local)?; + } Ok(registry) } @@ -204,12 +211,20 @@ impl AdapterRegistry { if !is_adapter_descriptor_file(&path) { continue; } - let descriptor = load_adapter_descriptor(&path)?; - self.register_descriptor(path, source, descriptor)?; + self.register_descriptor_file(&path, source)?; } Ok(()) } + fn register_descriptor_file( + &mut self, + path: &Path, + source: AdapterDescriptorSource, + ) -> Result<()> { + let descriptor = load_adapter_descriptor(path)?; + self.register_descriptor(path.to_path_buf(), source, descriptor) + } + fn register_descriptor( &mut self, path: PathBuf, @@ -989,6 +1004,16 @@ fn validate_config(config: &FabricConfig) -> Result<()> { pub fn resolve_run_plan_from_config( config: FabricConfig, context: ResolveContext, +) -> Result { + resolve_run_plan_from_config_with_adapter_descriptors(config, context, &[]) +} + +/// Resolve a typed Fabric config with caller-registered adapter descriptors. +#[doc(hidden)] +pub fn resolve_run_plan_from_config_with_adapter_descriptors( + config: FabricConfig, + context: ResolveContext, + adapter_descriptors: &[PathBuf], ) -> Result { validate_config(&config)?; let supplied_base_dir = context.base_dir; @@ -998,7 +1023,7 @@ pub fn resolve_run_plan_from_config( path: supplied_base_dir, source, })?; - resolve_run_plan(config, base_dir) + resolve_run_plan(config, base_dir, adapter_descriptors) } fn read_json(path: &Path) -> Result @@ -1015,8 +1040,12 @@ where }) } -fn resolve_run_plan(config: FabricConfig, base_dir: PathBuf) -> Result { - let adapter_descriptor = resolve_adapter_descriptor(&config, &base_dir)?; +fn resolve_run_plan( + config: FabricConfig, + base_dir: PathBuf, + adapter_descriptors: &[PathBuf], +) -> Result { + let adapter_descriptor = resolve_adapter_descriptor(&config, &base_dir, adapter_descriptors)?; let descriptor = adapter_descriptor .as_ref() .map(|adapter| &adapter.descriptor); @@ -1042,9 +1071,10 @@ fn resolve_run_plan(config: FabricConfig, base_dir: PathBuf) -> Result fn resolve_adapter_descriptor( config: &FabricConfig, base_dir: &Path, + adapter_descriptors: &[PathBuf], ) -> Result> { let adapter_id = &config.harness.adapter_id; - let registry = AdapterRegistry::from_config(config, base_dir)?; + let registry = AdapterRegistry::from_config(config, base_dir, adapter_descriptors)?; let Some(entry) = registry.get(adapter_id) else { return Err(FabricError::UnknownAdapter { adapter_id: adapter_id.clone(), @@ -1712,4 +1742,24 @@ mod tests { assert_eq!(descriptor.contract_version, ADAPTER_CONTRACT_VERSION); assert_eq!(descriptor.adapter_kind, AdapterKind::Python); } + + #[test] + fn resolves_caller_registered_adapter_descriptor() { + let descriptor_path = repository_root() + .join("tests/fixtures/hermes-shim-agent/adapters/hermes-shim/fabric-adapter.json"); + let plan = resolve_run_plan_from_config_with_adapter_descriptors( + typed_config("test.fabric.hermes_shim"), + ResolveContext::new("/tmp/fabric-base"), + std::slice::from_ref(&descriptor_path), + ) + .expect("registered adapter descriptor"); + let resolved = plan.adapter_descriptor.expect("adapter descriptor"); + + assert_eq!(resolved.descriptor.adapter_id, "test.fabric.hermes_shim"); + assert_eq!(resolved.source, AdapterDescriptorSource::Local); + assert_eq!( + resolved.path, + descriptor_path.canonicalize().expect("descriptor path") + ); + } } diff --git a/crates/fabric-python/src/lib.rs b/crates/fabric-python/src/lib.rs index 1362f4803..d6130ab32 100644 --- a/crates/fabric-python/src/lib.rs +++ b/crates/fabric-python/src/lib.rs @@ -5,9 +5,9 @@ use std::path::PathBuf; +use nemo_fabric_core::config::resolve_run_plan_from_config_with_adapter_descriptors; use nemo_fabric_core::{ - FabricConfig, ResolveContext, RunPlan, RunRequest, RuntimeHandle, doctor_plan, - resolve_run_plan_from_config, run_plan, + FabricConfig, ResolveContext, RunPlan, RunRequest, RuntimeHandle, doctor_plan, run_plan, }; use pyo3::exceptions::PyRuntimeError; use pyo3::prelude::*; @@ -20,33 +20,39 @@ fn version() -> PyResult { /// Resolve typed config JSON into a runnable plan and return JSON. #[pyfunction] -#[pyo3(signature = (config_json, base_dir=None))] -fn plan_config(py: Python<'_>, config_json: String, base_dir: Option) -> PyResult { +#[pyo3(signature = (config_json, base_dir=None, adapter_descriptors=None))] +fn plan_config( + py: Python<'_>, + config_json: String, + base_dir: Option, + adapter_descriptors: Option>, +) -> PyResult { let config = parse_config(config_json)?; let plan = py - .detach(|| resolve_run_plan_from_config(config, resolve_context(base_dir))) + .detach(|| resolve_config(config, base_dir, adapter_descriptors)) .map_err(to_py_error)?; to_json(&plan) } /// Diagnose typed config JSON without installing or running it. #[pyfunction] -#[pyo3(signature = (config_json, base_dir=None))] +#[pyo3(signature = (config_json, base_dir=None, adapter_descriptors=None))] fn doctor_config( py: Python<'_>, config_json: String, base_dir: Option, + adapter_descriptors: Option>, ) -> PyResult { let config = parse_config(config_json)?; let plan = py - .detach(|| resolve_run_plan_from_config(config, resolve_context(base_dir))) + .detach(|| resolve_config(config, base_dir, adapter_descriptors)) .map_err(to_py_error)?; to_json(&doctor_plan(&plan)) } /// Run typed config JSON through its Fabric adapter and return JSON. #[pyfunction] -#[pyo3(signature = (config_json, base_dir=None, input_text=None, input_file=None, request_json=None, request_file=None))] +#[pyo3(signature = (config_json, base_dir=None, input_text=None, input_file=None, request_json=None, request_file=None, adapter_descriptors=None))] fn run_config( py: Python<'_>, config_json: String, @@ -55,10 +61,11 @@ fn run_config( input_file: Option, request_json: Option, request_file: Option, + adapter_descriptors: Option>, ) -> PyResult { let config = parse_config(config_json)?; let plan = py - .detach(|| resolve_run_plan_from_config(config, resolve_context(base_dir))) + .detach(|| resolve_config(config, base_dir, adapter_descriptors)) .map_err(to_py_error)?; let request = match (request_file, request_json, input_file, input_text) { (Some(path), None, None, None) => std::fs::read_to_string(PathBuf::from(&path)) @@ -150,6 +157,23 @@ fn resolve_context(base_dir: Option) -> ResolveContext { ResolveContext::new(base_dir.unwrap_or_else(|| ".".to_string())) } +fn resolve_config( + config: FabricConfig, + base_dir: Option, + adapter_descriptors: Option>, +) -> nemo_fabric_core::Result { + let adapter_descriptors = adapter_descriptors + .unwrap_or_default() + .into_iter() + .map(PathBuf::from) + .collect::>(); + resolve_run_plan_from_config_with_adapter_descriptors( + config, + resolve_context(base_dir), + &adapter_descriptors, + ) +} + fn parse_config(contents: String) -> PyResult { serde_json::from_str(&contents).map_err(|error| PyRuntimeError::new_err(error.to_string())) } diff --git a/python/src/nemo_fabric/_native.pyi b/python/src/nemo_fabric/_native.pyi index 8df4b4bc0..2f538a873 100644 --- a/python/src/nemo_fabric/_native.pyi +++ b/python/src/nemo_fabric/_native.pyi @@ -5,10 +5,12 @@ def version() -> str: ... def plan_config( config_json: str, base_dir: str | None = None, + adapter_descriptors: list[str] | None = None, ) -> str: ... def doctor_config( config_json: str, base_dir: str | None = None, + adapter_descriptors: list[str] | None = None, ) -> str: ... def run_config( config_json: str, @@ -17,6 +19,7 @@ def run_config( input_file: str | None = None, request_json: str | None = None, request_file: str | None = None, + adapter_descriptors: list[str] | None = None, ) -> str: ... def start_runtime(plan_json: str) -> str: ... def invoke_runtime( diff --git a/python/src/nemo_fabric/client.py b/python/src/nemo_fabric/client.py index f6d23c10a..ec96b6092 100644 --- a/python/src/nemo_fabric/client.py +++ b/python/src/nemo_fabric/client.py @@ -7,9 +7,11 @@ import asyncio import importlib +import importlib.metadata import json import os from collections.abc import Mapping +from pathlib import Path from typing import Any from nemo_fabric.errors import ( FabricConfigError, @@ -86,6 +88,7 @@ def plan( raw = native.plan_config( _config_json(config), _base_dir_arg(base_dir), + _installed_adapter_descriptor_paths(), ) return RunPlan.from_mapping(json.loads(raw)) except FabricError: @@ -124,6 +127,7 @@ def diagnose() -> DoctorReport: raw = native.doctor_config( _config_json(config), _base_dir_arg(base_dir), + _installed_adapter_descriptor_paths(), ) return DoctorReport.from_mapping(json.loads(raw)) @@ -271,3 +275,15 @@ def _config_json(config: FabricConfig) -> str: def _base_dir_arg(base_dir: str | os.PathLike[str] | None) -> str | None: return None if base_dir is None else os.fspath(base_dir) + + +def _installed_adapter_descriptor_paths() -> list[str]: + paths: set[str] = set() + for distribution in importlib.metadata.distributions(): + for file in distribution.files or (): + if file.name != "fabric-adapter.json": + continue + path = Path(distribution.locate_file(file)) + if path.is_file(): + paths.add(os.fspath(path.resolve())) + return sorted(paths) diff --git a/tests/python/test_runtime.py b/tests/python/test_runtime.py index 582c351d8..2667ae7d1 100644 --- a/tests/python/test_runtime.py +++ b/tests/python/test_runtime.py @@ -95,7 +95,9 @@ async def _wait_for(event: threading.Event, timeout: float = 2.0) -> bool: def mock_native_fixture() -> MagicMock: mock_native = MagicMock() mock_native.requests = [] - mock_native.plan_config.side_effect = lambda config_json, base_dir: json.dumps(_plan()) + mock_native.plan_config.side_effect = ( + lambda config_json, base_dir, adapter_descriptors: json.dumps(_plan()) + ) mock_native.start_runtime.return_value = json.dumps(_runtime()) def invoke(plan_json: str, runtime_json: str, request_json: str) -> str: diff --git a/tests/python/test_sdk_contract.py b/tests/python/test_sdk_contract.py index 2b44272c5..4694e8200 100644 --- a/tests/python/test_sdk_contract.py +++ b/tests/python/test_sdk_contract.py @@ -10,8 +10,10 @@ from pathlib import Path from typing import Any from typing import get_overloads +from unittest.mock import MagicMock import nemo_fabric +import nemo_fabric.client as client_mod import nemo_fabric.errors as fabric_errors import pytest from nemo_fabric import AdapterInfo @@ -627,6 +629,7 @@ class NativeRecorder: def __init__(self) -> None: self.requests: list[dict[str, Any]] = [] self.config_base_dir_calls: list[str | None] = [] + self.adapter_descriptor_calls: list[list[str]] = [] self.stopped = 0 self.fail_invoke = False @@ -634,9 +637,11 @@ def plan_config( self, config_json: str, base_dir: str | None = None, + adapter_descriptors: list[str] | None = None, ) -> str: assert json.loads(config_json)["metadata"]["name"] == "demo" self.config_base_dir_calls.append(base_dir) + self.adapter_descriptor_calls.append(adapter_descriptors or []) return json.dumps(_plan()) def start_runtime(self, plan_json: str) -> str: @@ -1141,6 +1146,40 @@ def model_dump(self, *, mode: str, exclude_none: bool) -> dict[str, Any]: assert native.config_base_dir_calls == ["."] +def test_plan_registers_installed_adapter_descriptors( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +): + descriptor = ( + tmp_path + / "share" + / "nemo-fabric" + / "adapters" + / "shim" + / "fabric-adapter.json" + ) + descriptor.parent.mkdir(parents=True) + descriptor.write_text("{}") + + mock_distribution = MagicMock() + mock_distribution.files = [ + Path("share/nemo-fabric/adapters/shim/fabric-adapter.json") + ] + mock_distribution.locate_file.side_effect = lambda file: tmp_path / file + + monkeypatch.setattr( + client_mod.importlib.metadata, + "distributions", + lambda: [mock_distribution], + ) + native = NativeRecorder() + client = NativeClient(native) + + client.plan(_fabric_config()) + + assert native.adapter_descriptor_calls == [[str(descriptor.resolve())]] + + def test_fabric_config_constructors_emit_schema_shaped_mappings(): config = FabricConfig( metadata=MetadataConfig(name="demo"),