Skip to content

Commit

Permalink
Started on new set compute package
Browse files Browse the repository at this point in the history
  • Loading branch information
niklastheman committed Dec 1, 2023
1 parent 2b60de1 commit 5afc9d0
Show file tree
Hide file tree
Showing 2 changed files with 91 additions and 27 deletions.
83 changes: 75 additions & 8 deletions fedn/fedn/network/api/interface.py
Original file line number Diff line number Diff line change
Expand Up @@ -55,14 +55,16 @@ def _allowed_file_extension(
:param filename: The filename to check.
:type filename: str
:return: True if file extension is allowed, else False.
:rtype: bool
"""
:return: True and extension str if file extension is allowed, else False and None.
:rtype: Tuple (bool, str)
"""
if "." in filename:
extension = filename.rsplit(".", 1)[1].lower()
if extension in ALLOWED_EXTENSIONS:
return (True, extension)

return (False, None)

return (
"." in filename
and filename.rsplit(".", 1)[1].lower() in ALLOWED_EXTENSIONS
)

def get_clients(self, limit=None, skip=None, status=False):
"""Get all clients from the statestore.
Expand Down Expand Up @@ -207,7 +209,7 @@ def get_session(self, session_id):
payload[id] = info
return jsonify(payload)

def set_compute_package(self, file, helper_type):
def set_compute_package(self, file, helper_type: str):
"""Set the compute package in the statestore.
:param file: The compute package to set.
Expand All @@ -216,7 +218,72 @@ def set_compute_package(self, file, helper_type):
:rtype: :class:`flask.Response`
"""

if (
self.control.state() == ReducerState.instructing
or self.control.state() == ReducerState.monitoring
):
return (
jsonify(
{
"success": False,
"message": "Reducer is in instructing or monitoring state."
"Cannot set compute package.",
}
),
400,
)

if file is None:
return (
jsonify(
{
"success": False,
"message": "No file provided.",
}
),
404,
)

success, extension = self._allowed_file_extension(file.filename)

if not success:
return (
jsonify(
{
"success": False,
"message": f"File extension {extension} not allowed.",
}
),
404,
)

file_name = file.filename
storage_file_name = secure_filename(f"{str(uuid.uuid4())}.{extension}")

file_path = os.path.join("/app/client/package/", storage_file_name)
file.save(file_path)

self.control.set_compute_package(storage_file_name, file_path)
self.statestore.set_helper(helper_type)

success = self.statestore.set_compute_package(file_name, storage_file_name)

if not success:
return (
jsonify(
{
"success": False,
"message": "Failed to set compute package.",
}
),
400,
)

return jsonify({"success": True, "message": "Compute package set."})


if file and self._allowed_file_extension(file.filename):

filename = secure_filename(file.filename)
# TODO: make configurable, perhaps in config.py or package.py
file_path = os.path.join("/app/client/package/", filename)
Expand Down
35 changes: 16 additions & 19 deletions fedn/fedn/network/statestore/mongostatestore.py
Original file line number Diff line number Diff line change
Expand Up @@ -269,34 +269,31 @@ def get_validations(self, **kwargs):
result = self.control.validations.find(kwargs)
return result

def set_compute_package(self, filename):
def set_compute_package(self, file_name: str, storage_file_name: str):
"""Set the active compute package in statestore.
:param filename: The filename of the compute package.
:type filename: str
:param file_name: The file_name of the compute package.
:type file_name: str
:return: True if successful.
:rtype: bool
"""

obj = {
"file_name": file_name,
"storage_file_name": storage_file_name,
"committed_at": str(datetime.now()),
}

self.control.package.update_one(
{"key": "active"},
{
"$set": {
"filename": filename,
"committed_at": str(datetime.now()),
}
},
True,
)
self.control.package.update_one(
{"key": "package_trail"},
{
"$push": {
"filename": filename,
"committed_at": str(datetime.now()),
}
},
obj,
True,
)

trail_obj = {**{"key": "package_trail"}, **obj}

self.control.package.insert_one(trail_obj)

return True

def get_compute_package(self):
Expand Down

0 comments on commit 5afc9d0

Please sign in to comment.