623 lines
20 KiB
Python
623 lines
20 KiB
Python
"""Security and version-pinning tests for the installed SDK registry."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import replace
|
|
from importlib import metadata
|
|
from pathlib import Path
|
|
import py_compile
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from scimesh.sdk import (
|
|
AllowedPackage,
|
|
ArtifactCollection,
|
|
ArtifactRef,
|
|
CompatibilityError,
|
|
FeatureRequirement,
|
|
JobRequest,
|
|
LocalArtifactStore,
|
|
LocalPlanningContext,
|
|
PackageSpec,
|
|
TrustMode,
|
|
VersionRange,
|
|
WorkloadDefinition,
|
|
WorkloadId,
|
|
WorkloadRegistry,
|
|
installed_distribution_digest,
|
|
)
|
|
from scimesh.sdk.schema import (
|
|
ParameterValidationError,
|
|
validate_parameter_instance,
|
|
validate_schema_definition,
|
|
)
|
|
from scimesh.workloads.environment import current_scimesh_package_digest
|
|
from scimesh.workloads.library import default_sdk_runtime
|
|
from scimesh.workloads.search import similarity_search_sdk_definition
|
|
|
|
|
|
def _definition(
|
|
*, version: str = "1.0.0", digest_character: str = "a"
|
|
) -> WorkloadDefinition:
|
|
original = similarity_search_sdk_definition(shard_rows=2).definition()
|
|
manifest = replace(
|
|
original.manifest,
|
|
workload=WorkloadId("similarity-search", version),
|
|
package=PackageSpec("scimesh", "sha256:" + digest_character * 64),
|
|
)
|
|
return WorkloadDefinition(
|
|
manifest,
|
|
original.planner,
|
|
original.runners,
|
|
original.reducers,
|
|
original.verifiers,
|
|
)
|
|
|
|
|
|
def test_registry_requires_an_explicit_enabled_version_and_digest() -> None:
|
|
first = _definition(version="1.0.0", digest_character="a")
|
|
second = _definition(version="2.0.0", digest_character="b")
|
|
registry = WorkloadRegistry()
|
|
registry.register(first, enabled=True)
|
|
registry.register(second)
|
|
|
|
resolved, _ = registry.require("similarity-search", "1.0.0", "sha256:" + "a" * 64)
|
|
assert resolved is first
|
|
with pytest.raises(ValueError, match="unknown workload version"):
|
|
registry.require("similarity-search", "3.0.0", "sha256:" + "a" * 64)
|
|
with pytest.raises(ValueError, match="not enabled"):
|
|
registry.require("similarity-search", "2.0.0", "sha256:" + "b" * 64)
|
|
with pytest.raises(ValueError, match="not enabled"):
|
|
registry.require("similarity-search", "1.0.0", "sha256:" + "c" * 64)
|
|
with pytest.raises(ValueError, match="already registered"):
|
|
registry.register(first)
|
|
|
|
registry.enable("similarity-search", "2.0.0", "sha256:" + "b" * 64)
|
|
assert [item.workload.version for item in registry.descriptions()] == [
|
|
"1.0.0",
|
|
"2.0.0",
|
|
]
|
|
|
|
|
|
def test_compatibility_failure_occurs_before_planner_invocation(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
original = similarity_search_sdk_definition(shard_rows=2).definition()
|
|
|
|
class CountingPlanner:
|
|
calls = 0
|
|
entry_point = "tests.sdk_fixture:plan@v1"
|
|
|
|
def validate(self, request):
|
|
self.calls += 1
|
|
return original.planner.validate(request)
|
|
|
|
def plan(self, job, context):
|
|
self.calls += 1
|
|
return original.planner.plan(job, context)
|
|
|
|
planner = CountingPlanner()
|
|
definition = WorkloadDefinition(
|
|
original.manifest,
|
|
planner,
|
|
original.runners,
|
|
original.reducers,
|
|
original.verifiers,
|
|
)
|
|
registry = WorkloadRegistry()
|
|
registry.register(definition, enabled=True)
|
|
input_port = definition.manifest.inputs["input"]
|
|
artifact = ArtifactRef(
|
|
"11111111-1111-4111-8111-111111111111",
|
|
"a" * 64,
|
|
input_port.schema.ref,
|
|
input_port.schema.media_type,
|
|
1,
|
|
)
|
|
request = JobRequest(
|
|
definition.manifest.workload,
|
|
{"query_smiles": "CCO"},
|
|
{"input": ArtifactCollection.single(artifact)},
|
|
)
|
|
incompatible = replace(default_sdk_runtime(), protocol_version="2.0.0")
|
|
store = LocalArtifactStore(tmp_path / "artifacts")
|
|
|
|
with pytest.raises(CompatibilityError) as raised:
|
|
registry.plan(
|
|
request,
|
|
definition.manifest.package.digest,
|
|
incompatible,
|
|
LocalPlanningContext(store, store, tmp_path / "plan"),
|
|
)
|
|
assert raised.value.code == "protocol-mismatch"
|
|
assert planner.calls == 0
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("request_changes", "error_code"),
|
|
(
|
|
({"required_features": ("undeclared-feature",)}, "feature-undeclared"),
|
|
({"trust_mode": TrustMode.VERIFIED}, "trust-mode-undeclared"),
|
|
),
|
|
)
|
|
def test_job_selected_features_and_trust_mode_fail_closed_before_planning(
|
|
tmp_path: Path,
|
|
request_changes: dict[str, object],
|
|
error_code: str,
|
|
) -> None:
|
|
definition = similarity_search_sdk_definition(shard_rows=2).definition()
|
|
registry = WorkloadRegistry()
|
|
registry.register(definition, enabled=True)
|
|
input_port = definition.manifest.inputs["input"]
|
|
artifact = ArtifactRef(
|
|
"11111111-1111-4111-8111-111111111111",
|
|
"a" * 64,
|
|
input_port.schema.ref,
|
|
input_port.schema.media_type,
|
|
1,
|
|
records=1,
|
|
)
|
|
values: dict[str, object] = {
|
|
"workload": definition.manifest.workload,
|
|
"parameters": {"query_smiles": "CCO"},
|
|
"inputs": {"input": ArtifactCollection.single(artifact)},
|
|
}
|
|
values.update(request_changes)
|
|
request = JobRequest(**values) # type: ignore[arg-type]
|
|
store = LocalArtifactStore(tmp_path / "artifacts")
|
|
|
|
with pytest.raises(CompatibilityError) as raised:
|
|
registry.plan(
|
|
request,
|
|
definition.manifest.package.digest,
|
|
default_sdk_runtime(),
|
|
LocalPlanningContext(store, store, tmp_path / "plan"),
|
|
)
|
|
assert raised.value.code == error_code
|
|
|
|
|
|
class _EntryPoints(tuple):
|
|
def select(self, *, group: str):
|
|
assert group == WorkloadRegistry.ENTRY_POINT_GROUP
|
|
return self
|
|
|
|
|
|
def test_discovery_imports_only_an_exact_allowlisted_installed_entry_point(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
definition = similarity_search_sdk_definition().definition()
|
|
loaded: list[str] = []
|
|
|
|
class EntryPoint:
|
|
def __init__(self, name: str, distribution: str) -> None:
|
|
self.name = name
|
|
self.dist = (
|
|
metadata.distribution("scimesh")
|
|
if distribution == "scimesh"
|
|
else SimpleNamespace(name=distribution)
|
|
)
|
|
self.value = "scimesh.workloads.search:similarity_search_sdk_definition"
|
|
|
|
@property
|
|
def module(self) -> str:
|
|
return self.value.partition(":")[0]
|
|
|
|
def load(self):
|
|
loaded.append(self.name)
|
|
return lambda: definition
|
|
|
|
monkeypatch.setattr(
|
|
"scimesh.sdk.registry.metadata.entry_points",
|
|
lambda: _EntryPoints(
|
|
(
|
|
EntryPoint("evil-workload@1.0.0", "unapproved"),
|
|
EntryPoint("similarity-search@1.0.0", "scimesh"),
|
|
)
|
|
),
|
|
)
|
|
monkeypatch.setattr(
|
|
"scimesh.sdk.registry.installed_distribution_digest",
|
|
lambda _distribution: definition.manifest.package.digest,
|
|
)
|
|
registry = WorkloadRegistry()
|
|
registry.discover_installed(
|
|
(
|
|
AllowedPackage(
|
|
"scimesh",
|
|
definition.manifest.workload,
|
|
definition.manifest.package.digest,
|
|
),
|
|
)
|
|
)
|
|
|
|
assert loaded == ["similarity-search@1.0.0"]
|
|
assert registry.descriptions()[0].enabled
|
|
|
|
|
|
def test_discovery_measures_package_before_importing_entry_point(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
definition = similarity_search_sdk_definition().definition()
|
|
loaded = False
|
|
|
|
class EntryPoint:
|
|
name = "similarity-search@1.0.0"
|
|
dist = metadata.distribution("scimesh")
|
|
value = "scimesh.workloads.search:similarity_search_sdk_definition"
|
|
module = "scimesh.workloads.search"
|
|
|
|
def load(self):
|
|
nonlocal loaded
|
|
loaded = True
|
|
return lambda: definition
|
|
|
|
monkeypatch.setattr(
|
|
"scimesh.sdk.registry.metadata.entry_points",
|
|
lambda: _EntryPoints((EntryPoint(),)),
|
|
)
|
|
monkeypatch.setattr(
|
|
"scimesh.sdk.registry.installed_distribution_digest",
|
|
lambda _distribution: "sha256:" + "f" * 64,
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="content does not match"):
|
|
WorkloadRegistry().discover_installed(
|
|
(
|
|
AllowedPackage(
|
|
"scimesh",
|
|
definition.manifest.workload,
|
|
definition.manifest.package.digest,
|
|
),
|
|
)
|
|
)
|
|
assert loaded is False
|
|
|
|
|
|
def test_installed_digest_is_stable_when_python_generates_a_pycache(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
package = tmp_path / "fixture_pkg"
|
|
package.mkdir()
|
|
source = package / "__init__.py"
|
|
source.write_text("VALUE = 1\n", encoding="utf-8")
|
|
|
|
class FixtureDistribution:
|
|
name = "fixture-dist"
|
|
files = (Path("fixture_pkg/__init__.py"),)
|
|
entry_points = ()
|
|
|
|
@staticmethod
|
|
def read_text(name: str) -> str | None:
|
|
return "fixture_pkg\n" if name == "top_level.txt" else None
|
|
|
|
@staticmethod
|
|
def locate_file(value: object) -> Path:
|
|
return tmp_path / str(value)
|
|
|
|
distribution = FixtureDistribution()
|
|
before = installed_distribution_digest(distribution) # type: ignore[arg-type]
|
|
py_compile.compile(str(source), doraise=True)
|
|
|
|
assert installed_distribution_digest(distribution) == before # type: ignore[arg-type]
|
|
|
|
|
|
def test_discovery_rejects_entry_point_module_owned_by_another_distribution(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
owned = tmp_path / "owned_pkg"
|
|
owned.mkdir()
|
|
(owned / "__init__.py").write_text("", encoding="utf-8")
|
|
loaded = False
|
|
|
|
class Distribution:
|
|
name = "allowed-dist"
|
|
files = (Path("owned_pkg/__init__.py"),)
|
|
entry_points = ()
|
|
|
|
@staticmethod
|
|
def read_text(name: str) -> str | None:
|
|
return "owned_pkg\n" if name == "top_level.txt" else None
|
|
|
|
@staticmethod
|
|
def locate_file(value: object) -> Path:
|
|
return tmp_path / str(value)
|
|
|
|
class EntryPoint:
|
|
name = "similarity-search@1.0.0"
|
|
dist = Distribution()
|
|
value = "foreign_pkg.workload:factory"
|
|
module = "foreign_pkg.workload"
|
|
|
|
def load(self):
|
|
nonlocal loaded
|
|
loaded = True
|
|
raise AssertionError("foreign entry point must not load")
|
|
|
|
monkeypatch.setattr(
|
|
"scimesh.sdk.registry.metadata.entry_points",
|
|
lambda: _EntryPoints((EntryPoint(),)),
|
|
)
|
|
definition = similarity_search_sdk_definition().definition()
|
|
with pytest.raises(ValueError, match="outside its distribution"):
|
|
WorkloadRegistry().discover_installed(
|
|
(
|
|
AllowedPackage(
|
|
"allowed-dist",
|
|
definition.manifest.workload,
|
|
"sha256:" + "a" * 64,
|
|
),
|
|
)
|
|
)
|
|
assert loaded is False
|
|
|
|
|
|
def test_missing_allowlisted_entry_point_fails_without_loading_or_registering(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
loaded: list[str] = []
|
|
|
|
class EntryPoint:
|
|
name = "job-selected-module@1.0.0"
|
|
dist = SimpleNamespace(name="unapproved")
|
|
|
|
def load(self):
|
|
loaded.append(self.name)
|
|
raise AssertionError("unapproved entry point must not load")
|
|
|
|
monkeypatch.setattr(
|
|
"scimesh.sdk.registry.metadata.entry_points",
|
|
lambda: _EntryPoints((EntryPoint(),)),
|
|
)
|
|
registry = WorkloadRegistry()
|
|
with pytest.raises(ValueError, match="were not installed"):
|
|
registry.discover_installed(
|
|
(
|
|
AllowedPackage(
|
|
"scimesh",
|
|
WorkloadId("similarity-search", "1.0.0"),
|
|
current_scimesh_package_digest(),
|
|
),
|
|
)
|
|
)
|
|
assert loaded == []
|
|
assert registry.descriptions() == ()
|
|
|
|
|
|
def test_parameter_schema_accepts_finite_big_integer_bounds() -> None:
|
|
bound = 10**400
|
|
schema = {"type": "integer", "minimum": -bound, "maximum": bound}
|
|
|
|
validate_schema_definition(schema)
|
|
validate_parameter_instance(bound, schema)
|
|
|
|
with pytest.raises(ParameterValidationError, match="violates maximum"):
|
|
validate_parameter_instance(bound + 1, schema)
|
|
|
|
|
|
def test_job_parameters_reject_unbounded_json_integers_early() -> None:
|
|
with pytest.raises(ValueError, match="4096-bit JSON bound"):
|
|
JobRequest(
|
|
WorkloadId("similarity-search", "1.0.0"),
|
|
{"value": 10**2_000},
|
|
{},
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("value", "multiple", "accepted"),
|
|
[
|
|
(3 * 10**400, 3, True),
|
|
(10**400, 3, False),
|
|
(10**400, 0.1, True),
|
|
(0.3, 0.1, True),
|
|
(0.31, 0.1, False),
|
|
],
|
|
)
|
|
def test_parameter_schema_multiple_of_is_exact_without_float_overflow(
|
|
value: int | float,
|
|
multiple: int | float,
|
|
accepted: bool,
|
|
) -> None:
|
|
schema = {"type": "number", "multipleOf": multiple}
|
|
validate_schema_definition(schema)
|
|
|
|
if accepted:
|
|
validate_parameter_instance(value, schema)
|
|
else:
|
|
with pytest.raises(ParameterValidationError, match="violates multipleOf"):
|
|
validate_parameter_instance(value, schema)
|
|
|
|
|
|
def test_parameter_schema_equality_uses_json_types() -> None:
|
|
validate_schema_definition({"enum": [True, 1]})
|
|
with pytest.raises(ValueError, match="enum values must be unique"):
|
|
validate_schema_definition({"enum": [1, 1.0]})
|
|
|
|
validate_parameter_instance(True, {"enum": [True]})
|
|
with pytest.raises(ParameterValidationError, match="outside enum"):
|
|
validate_parameter_instance(1, {"enum": [True]})
|
|
validate_parameter_instance(1.0, {"enum": [1]})
|
|
|
|
validate_parameter_instance({"enabled": True}, {"const": {"enabled": True}})
|
|
with pytest.raises(ParameterValidationError, match="does not match const"):
|
|
validate_parameter_instance({"enabled": 1}, {"const": {"enabled": True}})
|
|
|
|
unique = {"type": "array", "uniqueItems": True}
|
|
validate_parameter_instance([True, 1, {"enabled": True}, {"enabled": 1}], unique)
|
|
with pytest.raises(ParameterValidationError, match="items must be unique"):
|
|
validate_parameter_instance([1, 1.0], unique)
|
|
|
|
|
|
def test_disabled_workload_is_not_resolvable_until_re_enabled() -> None:
|
|
definition = _definition()
|
|
registry = WorkloadRegistry()
|
|
registry.register(definition, enabled=True)
|
|
resolved, _ = registry.require("similarity-search", "1.0.0", "sha256:" + "a" * 64)
|
|
assert resolved is definition
|
|
|
|
registry.disable("similarity-search", "1.0.0", "sha256:" + "a" * 64)
|
|
with pytest.raises(ValueError, match="not enabled"):
|
|
registry.require("similarity-search", "1.0.0", "sha256:" + "a" * 64)
|
|
|
|
registry.enable("similarity-search", "1.0.0", "sha256:" + "a" * 64)
|
|
resolved, _ = registry.require("similarity-search", "1.0.0", "sha256:" + "a" * 64)
|
|
assert resolved is definition
|
|
|
|
|
|
def test_discovery_rechecks_the_package_digest_after_loading(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
definition = similarity_search_sdk_definition().definition()
|
|
digests = iter((definition.manifest.package.digest, "sha256:" + "e" * 64))
|
|
|
|
class EntryPoint:
|
|
name = "similarity-search@1.0.0"
|
|
dist = metadata.distribution("scimesh")
|
|
value = "scimesh.workloads.search:similarity_search_sdk_definition"
|
|
module = "scimesh.workloads.search"
|
|
|
|
def load(self):
|
|
return lambda: definition
|
|
|
|
monkeypatch.setattr(
|
|
"scimesh.sdk.registry.metadata.entry_points",
|
|
lambda: _EntryPoints((EntryPoint(),)),
|
|
)
|
|
monkeypatch.setattr(
|
|
"scimesh.sdk.registry.installed_distribution_digest",
|
|
lambda _distribution: next(digests),
|
|
)
|
|
|
|
registry = WorkloadRegistry()
|
|
with pytest.raises(ValueError, match="changed while loading"):
|
|
registry.discover_installed(
|
|
(
|
|
AllowedPackage(
|
|
"scimesh",
|
|
definition.manifest.workload,
|
|
definition.manifest.package.digest,
|
|
),
|
|
)
|
|
)
|
|
assert registry.descriptions() == ()
|
|
|
|
|
|
def test_request_trust_mode_must_be_enforceable_by_runtime_and_stages(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
original = similarity_search_sdk_definition(shard_rows=2).definition()
|
|
stages = tuple(
|
|
replace(stage, trust_modes=("trusted",))
|
|
for stage in original.manifest.workflow.stages
|
|
)
|
|
manifest = replace(
|
|
original.manifest,
|
|
workflow=replace(original.manifest.workflow, stages=stages),
|
|
trust_modes=(TrustMode.TRUSTED, TrustMode.VERIFIED),
|
|
)
|
|
definition = WorkloadDefinition(
|
|
manifest,
|
|
original.planner,
|
|
original.runners,
|
|
original.reducers,
|
|
original.verifiers,
|
|
)
|
|
registry = WorkloadRegistry()
|
|
registry.register(definition, enabled=True)
|
|
input_port = definition.manifest.inputs["input"]
|
|
artifact = ArtifactRef(
|
|
"11111111-1111-4111-8111-111111111111",
|
|
"a" * 64,
|
|
input_port.schema.ref,
|
|
input_port.schema.media_type,
|
|
1,
|
|
records=1,
|
|
)
|
|
request = JobRequest(
|
|
definition.manifest.workload,
|
|
{"query_smiles": "CCO"},
|
|
{"input": ArtifactCollection.single(artifact)},
|
|
trust_mode=TrustMode.VERIFIED,
|
|
)
|
|
store = LocalArtifactStore(tmp_path / "artifacts")
|
|
|
|
with pytest.raises(CompatibilityError) as raised:
|
|
registry.plan(
|
|
request,
|
|
manifest.package.digest,
|
|
default_sdk_runtime(),
|
|
LocalPlanningContext(store, store, tmp_path / "runtime-plan"),
|
|
)
|
|
assert raised.value.code == "trust-mode-unavailable"
|
|
|
|
runtime = replace(
|
|
default_sdk_runtime(),
|
|
trust_modes=(TrustMode.TRUSTED, TrustMode.VERIFIED),
|
|
)
|
|
with pytest.raises(CompatibilityError) as raised:
|
|
registry.plan(
|
|
request,
|
|
manifest.package.digest,
|
|
runtime,
|
|
LocalPlanningContext(store, store, tmp_path / "stage-plan"),
|
|
)
|
|
assert raised.value.code == "stage-trust-unavailable"
|
|
|
|
|
|
def test_job_cannot_require_a_feature_outside_the_runtime(tmp_path: Path) -> None:
|
|
original = similarity_search_sdk_definition(shard_rows=2).definition()
|
|
manifest = replace(
|
|
original.manifest,
|
|
optional_features=(
|
|
FeatureRequirement("gpu-fastpath", VersionRange(">=1,<2"), "cpu-fallback"),
|
|
),
|
|
)
|
|
definition = WorkloadDefinition(
|
|
manifest,
|
|
original.planner,
|
|
original.runners,
|
|
original.reducers,
|
|
original.verifiers,
|
|
)
|
|
registry = WorkloadRegistry()
|
|
registry.register(definition, enabled=True)
|
|
input_port = definition.manifest.inputs["input"]
|
|
artifact = ArtifactRef(
|
|
"11111111-1111-4111-8111-111111111111",
|
|
"a" * 64,
|
|
input_port.schema.ref,
|
|
input_port.schema.media_type,
|
|
1,
|
|
records=1,
|
|
)
|
|
request = JobRequest(
|
|
definition.manifest.workload,
|
|
{"query_smiles": "CCO"},
|
|
{"input": ArtifactCollection.single(artifact)},
|
|
required_features=("gpu-fastpath",),
|
|
)
|
|
|
|
negotiated = registry.require(
|
|
"similarity-search",
|
|
"1.0.0",
|
|
manifest.package.digest,
|
|
runtime=default_sdk_runtime(),
|
|
)[1]
|
|
assert negotiated is not None
|
|
assert negotiated.optional_fallbacks == {"gpu-fastpath": "cpu-fallback"}
|
|
|
|
with pytest.raises(CompatibilityError) as raised:
|
|
registry.plan(
|
|
request,
|
|
manifest.package.digest,
|
|
default_sdk_runtime(),
|
|
LocalPlanningContext(
|
|
LocalArtifactStore(tmp_path / "artifacts"),
|
|
LocalArtifactStore(tmp_path / "artifacts"),
|
|
tmp_path / "plan",
|
|
),
|
|
)
|
|
assert raised.value.code == "feature-unavailable"
|