diff --git a/fedn/network/api/v1/model_routes.py b/fedn/network/api/v1/model_routes.py index db52fae7e..9a3abdc47 100644 --- a/fedn/network/api/v1/model_routes.py +++ b/fedn/network/api/v1/model_routes.py @@ -736,3 +736,44 @@ def get_active_model(): return jsonify({"message": "No active model found"}), 404 except Exception: return jsonify({"message": "An unexpected error occurred"}), 500 + + +@bp.route("/active", methods=["PUT"]) +@jwt_auth_required(role="admin") +def set_active_model(): + """Set active model + Sets the active model (id). + --- + tags: + - Models + parameters: + - name: model + in: body + required: true + type: object + description: The model data to update + responses: + 200: + description: The updated active model id + schema: + type: string + 500: + description: An error occurred + schema: + type: object + properties: + message: + type: string + """ + try: + data = request.get_json() + model_id = data["id"] + + response = model_store.set_active(model_id) + + if response: + return jsonify({"message": "Active model set"}), 200 + else: + return jsonify({"message": "Failed to set active model"}), 500 + except Exception: + return jsonify({"message": "An unexpected error occurred"}), 500 diff --git a/fedn/network/api/v1/package_routes.py b/fedn/network/api/v1/package_routes.py index 0d559c2fa..d844a93ba 100644 --- a/fedn/network/api/v1/package_routes.py +++ b/fedn/network/api/v1/package_routes.py @@ -458,8 +458,12 @@ def set_active_package(): try: data = request.json package_id = data["id"] - package_store.set_active(package_id) - return jsonify({"message": "Active package set"}), 200 + response = package_store.set_active(package_id) + + if response: + return jsonify({"message": "Active package set"}), 200 + else: + return jsonify({"message": "Active package not set"}), 500 except Exception: return jsonify({"message": "An unexpected error occurred"}), 500 diff --git a/fedn/network/storage/statestore/stores/model_store.py b/fedn/network/storage/statestore/stores/model_store.py index c923f6d4a..27efcc9a3 100644 --- a/fedn/network/storage/statestore/stores/model_store.py +++ b/fedn/network/storage/statestore/stores/model_store.py @@ -26,7 +26,7 @@ def from_dict(data: dict) -> "Model": model=data["model"] if "model" in data else None, parent_model=data["parent_model"] if "parent_model" in data else None, session_id=data["session_id"] if "session_id" in data else None, - committed_at=data["committed_at"] if "committed_at" in data else None + committed_at=data["committed_at"] if "committed_at" in data else None, ) @@ -76,7 +76,7 @@ def update(self, id: str, item: Model) -> Tuple[bool, Any]: return super().update(id, item) - def add(self, item: Model)-> Tuple[bool, Any]: + def add(self, item: Model) -> Tuple[bool, Any]: raise NotImplementedError("Add not implemented for ModelStore") def delete(self, id: str) -> bool: @@ -104,10 +104,7 @@ def list(self, limit: int, skip: int, sort_key: str, sort_order=pymongo.DESCENDI response = super().list(limit, skip, sort_key or "committed_at", sort_order, use_typing=use_typing, **kwargs) result = [Model.from_dict(item) for item in response["result"]] if use_typing else response["result"] - return { - "count": response["count"], - "result": result - } + return {"count": response["count"], "result": result} def list_descendants(self, id: str, limit: int, use_typing: bool = False) -> List[Model]: """List descendants @@ -220,3 +217,26 @@ def get_active(self) -> str: raise EntityNotFound("Active model not found") return active_model["model"] + + def set_active(self, id: str) -> bool: + """Set the active model + param id: The id of the entity + type: str + description: The id of the entity, can be either the id or the model (property) + return: True if successful + """ + kwargs = {"key": "models"} + if ObjectId.is_valid(id): + id_obj = ObjectId(id) + kwargs["_id"] = id_obj + else: + kwargs["model"] = id + + model = self.database[self.collection].find_one(kwargs) + + if model is None: + raise EntityNotFound(f"Entity with (id | model) {id} not found") + + self.database[self.collection].update_one({"key": "current_model"}, {"$set": {"model": model["model"]}}) + + return True