# # This file is distributed under the MIT License. See LICENSE.md for details. # from dataclasses import dataclass, field from pathlib import Path from typing import Any from revng.pypeline.container import ContainerFormat from revng.pypeline.model import Model, ReadOnlyModel from revng.pypeline.object import Kind from revng.pypeline.pipeline import Pipeline from revng.pypeline.storage.storage_provider import FileStorageEntry from revng.pypeline.storage.storage_provider import storage_provider_factory_factory from revng.pypeline.task.requests import Requests from revng.pypeline.utils import bytes_to_string from revng.pypeline.utils.pipeline import deserialize_configuration, get_pipeline_description from revng.pypeline.utils.registry import get_singleton from .exceptions import EpochError, MalformedRequestError from .utils import compute_objects @dataclass class Response: """A simple wrapper used to enforce correctness and a standard format for responses.""" code: int """The HTTP status code of the response.""" body: Any """The response body of the request.""" content_type: str | None = None """The MIME of the response content.""" headers: dict[str, str] = field(default_factory=dict) """The headers of the response.""" notifications: list[Any] = field(default_factory=list) """List of notifications to send through the NotificationBroker.""" def to_dict(self): result = { "code": self.code, "body": self.body, } if len(self.headers) != 0: result["headers"] = self.headers if len(self.notifications) != 0: result["notifications"] = self.notifications return result class Daemon: """The transport agnostic part of the daemon.""" def __init__( self, pipeline: Pipeline, storage_provider_url: str, cache_dir: str, base_directory: Path, ): self.pipeline = pipeline self.cache_dir = cache_dir self.base_directory = base_directory self.storage_provider_factory = storage_provider_factory_factory(storage_provider_url) self.pipeline_description = get_pipeline_description(pipeline) def _get_storage_provider_context(self, request): project_id = request.get("project_id") token = request.get("token") return self.storage_provider_factory.get( base_directory=self.base_directory, project_id=project_id, token=token, cache_dir=self.cache_dir, ) async def get_epoch(self, request) -> Response: storage_provider_context = self._get_storage_provider_context(request) async with storage_provider_context as storage_provider: return Response(code=200, body={"epoch": storage_provider.get_epoch()}) async def get_model(self, request): storage_provider_context = self._get_storage_provider_context(request) async with storage_provider_context as storage_provider: model, epoch = storage_provider.get_model() return Response( code=200, body={ "epoch": epoch, "model": bytes_to_string(model.serialize(), is_text=model.__class__.is_text()), }, ) def get_pipeline(self) -> Response: return Response( code=200, body=self.pipeline_description, ) async def put_file(self, request) -> Response: entry = FileStorageEntry(request["name"], contents=request["contents"]) storage_provider_context = self._get_storage_provider_context(request) async with storage_provider_context as storage_provider: hashes = storage_provider.put_files_in_storage([entry]) return Response( code=200, body={ "name": request["name"], "hash": hashes[0], }, ) async def artifact(self, request) -> Response: artifact_name: str = request["artifact"] objects: list[str] | None = request.get("objects") raw_configuration: dict[str, str] = request.get("configuration", {}) epoch: int = request["epoch"] format_: str = request.get("format", "json") # Validate data if artifact_name not in self.pipeline.artifacts: raise MalformedRequestError(f"Artifact {artifact_name} not found in the pipeline") if format_ not in ("json", "tar"): raise MalformedRequestError(f"Format {format_} is not valid, valid values: json, tar") # Convert configuration configuration = deserialize_configuration(self.pipeline, raw_configuration) artifact = self.pipeline.artifacts[artifact_name] configuration_hash = artifact.node.configuration_id(configuration) # Compute the artifact storage_provider_context = self._get_storage_provider_context(request) async with storage_provider_context as storage_provider: # Load the model model, real_epoch = storage_provider.get_model() if real_epoch != epoch: raise EpochError(real_epoch, epoch) object_set = compute_objects( model=ReadOnlyModel(model), kind=artifact.container.container_type.kind, objects=objects, ) container = self.pipeline.get_artifact( model=ReadOnlyModel(model), artifact=artifact, requests=object_set, configuration=configuration, storage_provider=storage_provider, ) headers = {"x-pypeline-configuration-hash": configuration_hash} if format_ == "json": return Response( code=200, body={ key: bytes_to_string(value, container.is_text()) for key, value in container.to_dict(object_set).items() }, headers=headers, ) else: return Response( code=200, body=container.to_bytes(object_set, ContainerFormat.TAR), content_type="application/x-tar", headers=headers, ) async def analyze(self, request) -> Response: """Process analysis requests""" # Extract the data epoch = request["epoch"] analysis = request["analysis"] raw_configuration = request.get("configuration", {}) containers = request.get("containers", {}) # Validate data and normalize to analysis list if analysis not in self.pipeline.analyses and analysis not in self.pipeline.analysis_lists: raise MalformedRequestError(f"Analysis {analysis} not found in the pipeline") # Check that the given containers are declared in the pipeline for container_name, objects in containers.items(): for decl in self.pipeline.declarations: if container_name == decl.name: break else: raise MalformedRequestError(f"Container {container_name} not found in the pipeline") configuration = deserialize_configuration(self.pipeline, raw_configuration) storage_provider_context = self._get_storage_provider_context(request) async with storage_provider_context as storage_provider: # Load the model model, real_epoch = storage_provider.get_model() if real_epoch != epoch: raise EpochError(real_epoch, epoch) # Run an analysis if analysis in self.pipeline.analyses: current_model = ReadOnlyModel(model) # Setup the requests requests = Requests() for binding in self.pipeline.analyses[analysis].bindings: kind: Kind = binding.container_type.kind objects = containers.get(binding.name) if objects is not None and not isinstance(objects, list): raise MalformedRequestError( f"Objects for container {binding.name} must be a list, " f"got {type(objects)}" ) requests.insert(binding, compute_objects(current_model, kind, objects)) # Run the analysis new_model, invalidated = self.pipeline.run_analysis( model=current_model, analysis_name=analysis, requests=requests, configuration=configuration, storage_provider=storage_provider, ) else: analysis_list = self.pipeline.analysis_lists[analysis] new_model, invalidated = self.pipeline.run_analysis_list( model=ReadOnlyModel(model), analysis_list=analysis_list, configuration=configuration, storage_provider=storage_provider, ) # Compute the diff between the original model and the final one diff_raw = model.diff(new_model).serialize() # TODO: this can be done much more efficiently new_epoch = storage_provider.get_epoch() # Only return cacheable artifacts invalidations invalidated_artifacts: list[dict[str, Any]] = [] for container_location, object_ids in invalidated.items(): artifact = self.pipeline.savepoint_id_to_artifact.get(container_location.savepoint_id) if artifact is None: continue invalidated_artifacts.append( { "name": artifact.name, "configuration": container_location.configuration_id, "object_ids": [object_id.serialize() for object_id in object_ids], } ) model_type = get_singleton(Model) # type: ignore[type-abstract] diff = bytes_to_string(diff_raw, model_type.is_text()) # Return the updated model return Response( code=200, body={ "epoch": new_epoch, "diff": diff, }, notifications=[ { "type": "analysis", "analysis": analysis, "epoch": new_epoch, "diff": diff, "invalidated": invalidated_artifacts, } ], )