From 1defce19fb2822b3bc9463f79eb66a782fb6bae4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pablo=20Arag=C3=B3n?= Date: Tue, 2 Dec 2025 17:07:56 +0100 Subject: [PATCH 1/4] feat: change storages arguments to function instance --- gateway/api/ray.py | 8 ++------ gateway/api/serializers.py | 6 +----- gateway/api/services/arguments_storage.py | 14 +++++++++++--- gateway/api/services/file_storage.py | 16 ++++++++++++---- gateway/api/services/result_storage.py | 10 ++++++++-- gateway/api/use_cases/files/delete.py | 3 +-- gateway/api/use_cases/files/download.py | 3 +-- gateway/api/use_cases/files/list.py | 3 +-- gateway/api/use_cases/files/provider_delete.py | 3 +-- gateway/api/use_cases/files/provider_download.py | 3 +-- gateway/api/use_cases/files/provider_list.py | 3 +-- gateway/api/use_cases/files/provider_upload.py | 3 +-- gateway/api/use_cases/files/upload.py | 3 +-- gateway/api/use_cases/jobs/retrieve.py | 2 +- gateway/api/use_cases/jobs/save_result.py | 2 +- gateway/tests/api/test_v1_program.py | 8 ++++---- 16 files changed, 48 insertions(+), 42 deletions(-) diff --git a/gateway/api/ray.py b/gateway/api/ray.py index ad838c7bc..a0ea0a639 100644 --- a/gateway/api/ray.py +++ b/gateway/api/ray.py @@ -223,24 +223,20 @@ def _create_cluster_data(job: Job, cluster_name: str): # configure provider configuration if needed node_image = settings.RAY_NODE_IMAGE - provider_name = None if job.program.provider is not None: node_image = job.program.image - provider_name = job.program.provider.name user_file_storage = FileStorage( username=user.username, working_dir=WorkingDir.USER_STORAGE, - function_title=job.program.title, - provider_name=provider_name, + function=job.program, ) provider_file_storage = user_file_storage if job.program.provider is not None: provider_file_storage = FileStorage( username=user.username, working_dir=WorkingDir.PROVIDER_STORAGE, - function_title=job.program.title, - provider_name=provider_name, + function=job.program, ) cluster = get_template("rayclustertemplate.yaml") diff --git a/gateway/api/serializers.py b/gateway/api/serializers.py index d4193d2e1..d771be910 100644 --- a/gateway/api/serializers.py +++ b/gateway/api/serializers.py @@ -271,8 +271,6 @@ def create(self, validated_data): logger.info("Creating Job with RunExistingJobSerializer") status = Job.QUEUED program = validated_data.get("program") - function_title = validated_data.get("function_title") - provider_name = validated_data.get("provider_name", None) arguments = validated_data.get("arguments", "{}") author = validated_data.get("author") config = validated_data.get("config", None) @@ -301,9 +299,7 @@ def create(self, validated_data): ) ) - arguments_storage = ArgumentsStorage( - author.username, function_title, provider_name - ) + arguments_storage = ArgumentsStorage(program) arguments_storage.save(job.id, arguments) try: diff --git a/gateway/api/services/arguments_storage.py b/gateway/api/services/arguments_storage.py index c37acffc3..ac69ad308 100644 --- a/gateway/api/services/arguments_storage.py +++ b/gateway/api/services/arguments_storage.py @@ -15,9 +15,17 @@ class ArgumentsStorage: ARGUMENTS_FILE_EXTENSION = ".json" ENCODING = "utf-8" - def __init__( - self, username: str, function_title: str, provider_name: Optional[str] - ): + def __init__(self, function): + """ + Initialize ArgumentsStorage with a function instance. + + Args: + function: Program model instance containing title, provider, and author + """ + username = function.author.username + function_title = function.title + provider_name = function.provider.name if function.provider else None + # We need to use the same path as the FileStorage here # because it is attached the volume in the docker image if provider_name is None: diff --git a/gateway/api/services/file_storage.py b/gateway/api/services/file_storage.py index 6f2e18885..016420bd0 100644 --- a/gateway/api/services/file_storage.py +++ b/gateway/api/services/file_storage.py @@ -39,21 +39,29 @@ class FileStorage: Attributes: username (str): storage user's username working_dir (WorkingDir(Enum)): working directory - function_title (str): title of the function in case is needed to build the path - provider_name (str | None): name of the provider in caseis needed to build the path """ def __init__( self, username: str, working_dir: WorkingDir, - function_title: str, - provider_name: Optional[str], + function, ) -> None: + """ + Initialize FileStorage with a function instance. + + Args: + username: User's username + working_dir: Working directory type (USER_STORAGE or PROVIDER_STORAGE) + function: Program model instance containing title and provider + """ self.sub_path = None self.absolute_path = None self.username = username + function_title = function.title + provider_name = function.provider.name if function.provider else None + if working_dir is WorkingDir.USER_STORAGE: self.sub_path = self.__get_user_sub_path(function_title, provider_name) elif working_dir is WorkingDir.PROVIDER_STORAGE: diff --git a/gateway/api/services/result_storage.py b/gateway/api/services/result_storage.py index f33d22e8e..692ac8f01 100644 --- a/gateway/api/services/result_storage.py +++ b/gateway/api/services/result_storage.py @@ -15,8 +15,14 @@ class ResultStorage: RESULT_FILE_EXTENSION = ".json" ENCODING = "utf-8" - def __init__(self, username: str): - """Initialize the storage path for a given user.""" + def __init__(self, function): + """ + Initialize the storage path for a given function. + + Args: + function: Program model instance containing author + """ + username = function.author.username self.user_results_directory = os.path.join( settings.MEDIA_ROOT, username, "results" ) diff --git a/gateway/api/use_cases/files/delete.py b/gateway/api/use_cases/files/delete.py index c520970d2..092314a15 100644 --- a/gateway/api/use_cases/files/delete.py +++ b/gateway/api/use_cases/files/delete.py @@ -47,8 +47,7 @@ def execute( file_storage = FileStorage( username=user.username, working_dir=self.working_dir, - function_title=function_title, - provider_name=provider_name, + function=function, ) result = file_storage.remove_file(file_name=file_name) diff --git a/gateway/api/use_cases/files/download.py b/gateway/api/use_cases/files/download.py index a63ba02ba..ebdd53237 100644 --- a/gateway/api/use_cases/files/download.py +++ b/gateway/api/use_cases/files/download.py @@ -47,8 +47,7 @@ def execute( file_storage = FileStorage( username=user.username, working_dir=self.working_dir, - function_title=function_title, - provider_name=provider_name, + function=function, ) result = file_storage.get_file(file_name=requested_file_name) diff --git a/gateway/api/use_cases/files/list.py b/gateway/api/use_cases/files/list.py index 3c93a52b1..bfbd0e19d 100644 --- a/gateway/api/use_cases/files/list.py +++ b/gateway/api/use_cases/files/list.py @@ -41,8 +41,7 @@ def execute(self, user: AbstractUser, provider_name: str, function_title: str): file_storage = FileStorage( username=user.username, working_dir=self.working_dir, - function_title=function_title, - provider_name=provider_name, + function=function, ) return file_storage.get_files() diff --git a/gateway/api/use_cases/files/provider_delete.py b/gateway/api/use_cases/files/provider_delete.py index 7f748cfa6..b877e23d7 100644 --- a/gateway/api/use_cases/files/provider_delete.py +++ b/gateway/api/use_cases/files/provider_delete.py @@ -54,8 +54,7 @@ def execute( file_storage = FileStorage( username=user.username, working_dir=self.working_dir, - function_title=function_title, - provider_name=provider_name, + function=function, ) result = file_storage.remove_file(file_name=file_name) diff --git a/gateway/api/use_cases/files/provider_download.py b/gateway/api/use_cases/files/provider_download.py index 58806b406..ef44e116d 100644 --- a/gateway/api/use_cases/files/provider_download.py +++ b/gateway/api/use_cases/files/provider_download.py @@ -54,8 +54,7 @@ def execute( file_storage = FileStorage( username=user.username, working_dir=self.working_dir, - function_title=function_title, - provider_name=provider_name, + function=function, ) result = file_storage.get_file(file_name=requested_file_name) diff --git a/gateway/api/use_cases/files/provider_list.py b/gateway/api/use_cases/files/provider_list.py index 810bd5284..4d4ff4a4a 100644 --- a/gateway/api/use_cases/files/provider_list.py +++ b/gateway/api/use_cases/files/provider_list.py @@ -48,8 +48,7 @@ def execute(self, user: AbstractUser, provider_name: str, function_title: str): file_storage = FileStorage( username=user.username, working_dir=self.working_dir, - function_title=function_title, - provider_name=provider_name, + function=function, ) return file_storage.get_files() diff --git a/gateway/api/use_cases/files/provider_upload.py b/gateway/api/use_cases/files/provider_upload.py index c3d0f5f0b..bb7864cd8 100644 --- a/gateway/api/use_cases/files/provider_upload.py +++ b/gateway/api/use_cases/files/provider_upload.py @@ -56,8 +56,7 @@ def execute( file_storage = FileStorage( username=user.username, working_dir=self.working_dir, - function_title=function_title, - provider_name=provider_name, + function=function, ) result = file_storage.upload_file(file=uploaded_file) diff --git a/gateway/api/use_cases/files/upload.py b/gateway/api/use_cases/files/upload.py index 5b67afa8b..0dc156fd7 100644 --- a/gateway/api/use_cases/files/upload.py +++ b/gateway/api/use_cases/files/upload.py @@ -48,8 +48,7 @@ def execute( file_storage = FileStorage( username=user.username, working_dir=self.working_dir, - function_title=function_title, - provider_name=provider_name, + function=function, ) result = file_storage.upload_file(file=uploaded_file) diff --git a/gateway/api/use_cases/jobs/retrieve.py b/gateway/api/use_cases/jobs/retrieve.py index 5cdc01d36..c7b717552 100644 --- a/gateway/api/use_cases/jobs/retrieve.py +++ b/gateway/api/use_cases/jobs/retrieve.py @@ -38,7 +38,7 @@ def execute(self, job_id: UUID, user: AbstractUser, with_result: bool) -> Job: can_read_result = JobAccessPolicies.can_read_result(user, job) if with_result and can_read_result: - result_store = ResultStorage(user.username) + result_store = ResultStorage(job.program) result = result_store.get(str(job.id)) if result is not None: job.result = result diff --git a/gateway/api/use_cases/jobs/save_result.py b/gateway/api/use_cases/jobs/save_result.py index 214a0401f..824ae40e5 100644 --- a/gateway/api/use_cases/jobs/save_result.py +++ b/gateway/api/use_cases/jobs/save_result.py @@ -43,7 +43,7 @@ def execute(self, job_id: UUID, user: AbstractUser, result: str) -> Job: if not can_save_result: raise NotFoundError(f"Job [{job_id}] not found") - result_storage = ResultStorage(user.username) + result_storage = ResultStorage(job.program) result_storage.save(job.id, result) job.result = result diff --git a/gateway/tests/api/test_v1_program.py b/gateway/tests/api/test_v1_program.py index 20a3badbb..977ef0d53 100644 --- a/gateway/tests/api/test_v1_program.py +++ b/gateway/tests/api/test_v1_program.py @@ -146,7 +146,8 @@ def test_run(self): self.assertEqual(job.config.workers, None) self.assertEqual(job.config.auto_scaling, True) - arguments_storage = ArgumentsStorage(user.username, "Program", None) + program = Program.objects.get(title="Program") + arguments_storage = ArgumentsStorage(program) stored_arguments = arguments_storage.get(job.id) self.assertEqual(stored_arguments, arguments) @@ -195,9 +196,8 @@ def test_provider_run(self): self.assertEqual(job.config.workers, None) self.assertEqual(job.config.auto_scaling, True) - arguments_storage = ArgumentsStorage( - user.username, "Docker-Image-Program", "default" - ) + program = Program.objects.get(title="Docker-Image-Program") + arguments_storage = ArgumentsStorage(program) stored_arguments = arguments_storage.get(job.id) self.assertEqual(stored_arguments, arguments) From 203f6897ea92afb01c4198092f9bd2f2e6e7248b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pablo=20Arag=C3=B3n?= Date: Wed, 3 Dec 2025 10:55:37 +0100 Subject: [PATCH 2/4] fix tests --- gateway/tests/api/test_v1_program.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/gateway/tests/api/test_v1_program.py b/gateway/tests/api/test_v1_program.py index 977ef0d53..a55e705ee 100644 --- a/gateway/tests/api/test_v1_program.py +++ b/gateway/tests/api/test_v1_program.py @@ -146,7 +146,7 @@ def test_run(self): self.assertEqual(job.config.workers, None) self.assertEqual(job.config.auto_scaling, True) - program = Program.objects.get(title="Program") + program = Program.objects.get(title="Program", author=user) arguments_storage = ArgumentsStorage(program) stored_arguments = arguments_storage.get(job.id) @@ -196,7 +196,7 @@ def test_provider_run(self): self.assertEqual(job.config.workers, None) self.assertEqual(job.config.auto_scaling, True) - program = Program.objects.get(title="Docker-Image-Program") + program = Program.objects.get(title="Docker-Image-Program", author=user) arguments_storage = ArgumentsStorage(program) stored_arguments = arguments_storage.get(job.id) From 90410d36167ce44d621016a7b250c8e3582b97b1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pablo=20Arag=C3=B3n?= Date: Wed, 3 Dec 2025 11:56:02 +0100 Subject: [PATCH 3/4] add release notes --- .../notes/improve-storage-logic-cffa12b8e089a34d.yaml | 9 +++++++++ 1 file changed, 9 insertions(+) create mode 100644 releasenotes/notes/improve-storage-logic-cffa12b8e089a34d.yaml diff --git a/releasenotes/notes/improve-storage-logic-cffa12b8e089a34d.yaml b/releasenotes/notes/improve-storage-logic-cffa12b8e089a34d.yaml new file mode 100644 index 000000000..bf198af4e --- /dev/null +++ b/releasenotes/notes/improve-storage-logic-cffa12b8e089a34d.yaml @@ -0,0 +1,9 @@ +--- +upgrade: + - | + Refactored storage service initialization to accept ``Program`` model instances + instead of individual parameters. The ``FileStorage``, ``ArgumentsStorage``, and + ``ResultStorage`` classes now accept a ``function`` parameter (Program model instance) + rather than separate ``username``, ``function_title``, and ``provider_name`` parameters. + This change improves code cohesion by reducing the number of parameters passed between + components and ensures consistent access to program metadata across storage services. From a25c5e98d4703504d2a790f8b70bfc8ae6feec58 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pablo=20Arag=C3=B3n?= Date: Thu, 4 Dec 2025 12:24:43 +0100 Subject: [PATCH 4/4] fix: username from request, not function --- gateway/api/serializers.py | 2 +- gateway/api/services/arguments_storage.py | 6 ++++-- gateway/api/services/file_storage.py | 4 ---- gateway/tests/api/test_v1_program.py | 4 ++-- 4 files changed, 7 insertions(+), 9 deletions(-) diff --git a/gateway/api/serializers.py b/gateway/api/serializers.py index d771be910..832f24b54 100644 --- a/gateway/api/serializers.py +++ b/gateway/api/serializers.py @@ -299,7 +299,7 @@ def create(self, validated_data): ) ) - arguments_storage = ArgumentsStorage(program) + arguments_storage = ArgumentsStorage(author.username, program) arguments_storage.save(job.id, arguments) try: diff --git a/gateway/api/services/arguments_storage.py b/gateway/api/services/arguments_storage.py index ac69ad308..2e94c1652 100644 --- a/gateway/api/services/arguments_storage.py +++ b/gateway/api/services/arguments_storage.py @@ -6,6 +6,8 @@ from typing import Optional from django.conf import settings +from api.models import Program + logger = logging.getLogger("gateway") @@ -15,14 +17,14 @@ class ArgumentsStorage: ARGUMENTS_FILE_EXTENSION = ".json" ENCODING = "utf-8" - def __init__(self, function): + def __init__(self, username: str, function: Program): """ Initialize ArgumentsStorage with a function instance. Args: + username: Program model instance containing title, provider, and author function: Program model instance containing title, provider, and author """ - username = function.author.username function_title = function.title provider_name = function.provider.name if function.provider else None diff --git a/gateway/api/services/file_storage.py b/gateway/api/services/file_storage.py index 016420bd0..c80cb0c94 100644 --- a/gateway/api/services/file_storage.py +++ b/gateway/api/services/file_storage.py @@ -35,10 +35,6 @@ class WorkingDir(Enum): class FileStorage: """ The main objective of this class is to manage the access of the users to their storage. - - Attributes: - username (str): storage user's username - working_dir (WorkingDir(Enum)): working directory """ def __init__( diff --git a/gateway/tests/api/test_v1_program.py b/gateway/tests/api/test_v1_program.py index a55e705ee..10cb49165 100644 --- a/gateway/tests/api/test_v1_program.py +++ b/gateway/tests/api/test_v1_program.py @@ -147,7 +147,7 @@ def test_run(self): self.assertEqual(job.config.auto_scaling, True) program = Program.objects.get(title="Program", author=user) - arguments_storage = ArgumentsStorage(program) + arguments_storage = ArgumentsStorage(user.username, program) stored_arguments = arguments_storage.get(job.id) self.assertEqual(stored_arguments, arguments) @@ -197,7 +197,7 @@ def test_provider_run(self): self.assertEqual(job.config.auto_scaling, True) program = Program.objects.get(title="Docker-Image-Program", author=user) - arguments_storage = ArgumentsStorage(program) + arguments_storage = ArgumentsStorage(user.username, program) stored_arguments = arguments_storage.get(job.id) self.assertEqual(stored_arguments, arguments)