add arguments_schema field and preserve optional fields on program re-upload

Author: avilches

gateway/api/migrations/0055_program_arguments_schema.py

@@ -0,0 +1,18 @@
+# Generated by Django 5.2.12 on 2026-07-09 08:57
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("api", "0054_job_running_started_at"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="program",
+            name="arguments_schema",
+            field=models.TextField(blank=True, default="{}"),
+        ),
+    ]

gateway/api/use_cases/programs/upload.py

@@ -71,18 +71,21 @@ def execute(
     def _create(self, data: UploadFunctionInput, user, provider) -> Function:
         logger.info("user_id=%s program=%s | Creating function", user.id, data.title)
 
+        if data.entrypoint is None and data.image is None:
+            raise DRFValidationError("At least one of attributes (entrypoint, image) is required.")
+
         env_vars = data.env_vars
         if env_vars:
             env_vars = json.dumps(encrypt_env_vars(json.loads(env_vars)))
 
-        raw_deps = json.loads(data.dependencies)
+        raw_deps = json.loads(data.dependencies or "[]")
         dependencies = json.dumps([_normalize_dependency(d) for d in raw_deps])
 
         function = Function(
             title=data.title,
             author=user,
             provider=provider,
-            runner=data.runner,
+            runner=data.runner or Function.RAY,
             entrypoint=data.entrypoint or DEFAULT_PROGRAM_ENTRYPOINT,
             artifact=data.artifact,
             image=data.image,
@@ -93,6 +96,8 @@ def _create(self, data: UploadFunctionInput, user, provider) -> Function:
         )
         if data.type is not None:
             function.type = data.type
+        if data.arguments_schema is not None:
+            function.arguments_schema = data.arguments_schema
 
         CodeEngineProject.objects.assign_to_program(function)
         if function.runner == Function.FLEETS and not function.code_engine_project:
@@ -104,21 +109,23 @@ def _update(self, instance: Function, data: UploadFunctionInput, user) -> Functi
         logger.info("user_id=%s program=%s | Updating function", user.id, instance.title)
 
         instance.entrypoint = data.entrypoint or DEFAULT_PROGRAM_ENTRYPOINT
-        raw_deps = json.loads(data.dependencies)
-        instance.dependencies = json.dumps([_normalize_dependency(d) for d in raw_deps])
-        env_vars = data.env_vars
-        if env_vars:
-            env_vars = json.dumps(encrypt_env_vars(json.loads(env_vars)))
-        instance.env_vars = env_vars or {}
+        if data.dependencies is not None:
+            raw_deps = json.loads(data.dependencies)
+            instance.dependencies = json.dumps([_normalize_dependency(d) for d in raw_deps])
+        if data.env_vars is not None:
+            instance.env_vars = json.dumps(encrypt_env_vars(json.loads(data.env_vars)))
         instance.artifact = data.artifact
         instance.author = user
         instance.image = data.image
-        instance.runner = data.runner
+        if data.runner is not None:
+            instance.runner = data.runner
 
         if data.description is not None:
             instance.description = data.description
         if data.version is not None:
             instance.version = data.version
+        if data.arguments_schema is not None:
+            instance.arguments_schema = data.arguments_schema
 
         CodeEngineProject.objects.assign_to_program(instance)
         instance.save()

gateway/api/use_cases/programs/upload_input.py

@@ -3,8 +3,6 @@
 from dataclasses import dataclass
 from typing import Any
 
-from core.models import Program as Function
-
 
 def _parse_provider_and_title(provider_raw: str | None, title_raw: str) -> tuple[str | None, str]:
     if provider_raw:
@@ -28,11 +26,12 @@ class UploadFunctionInput:  # pylint: disable=too-many-instance-attributes
     artifact: Any = None
     image: str | None = None
     env_vars: str | None = None
-    dependencies: str = "[]"
-    runner: str = Function.RAY
+    dependencies: str | None = None
+    runner: str | None = None
     description: str | None = None
     version: str | None = None
     type: str | None = None
+    arguments_schema: str | None = None
 
     @classmethod
     def from_validated_data(cls, data: dict) -> "UploadFunctionInput":
@@ -45,9 +44,10 @@ def from_validated_data(cls, data: dict) -> "UploadFunctionInput":
             artifact=data.get("artifact"),
             image=data.get("image"),
             env_vars=data.get("env_vars"),
-            dependencies=data.get("dependencies", "[]"),
-            runner=data.get("runner", "ray"),
+            dependencies=data.get("dependencies"),
+            runner=data.get("runner"),
             description=data.get("description"),
             version=data.get("version"),
             type=data.get("type"),
+            arguments_schema=data.get("arguments_schema"),
         )

gateway/api/v1/views/programs/upload.py

@@ -35,6 +35,7 @@ class ProgramSerializer(serializers.ModelSerializer):
     image = serializers.CharField(required=False)
     provider = serializers.CharField(required=False)
     runner = serializers.CharField(required=False)
+    arguments_schema = serializers.CharField(required=False)
 
     class Meta:
         model = Program
@@ -50,6 +51,7 @@ class Meta:
             "type",
             "version",
             "runner",
+            "arguments_schema",
         ]
         ref_name = "ProgramsUploadProgram"
 
@@ -93,6 +95,14 @@ def validate_image(self, value):
         """Validate image."""
         return value
 
+    def validate_arguments_schema(self, value):
+        """Validates that arguments_schema is valid JSON."""
+        try:
+            json.loads(value)
+        except (json.JSONDecodeError, ValueError) as exc:
+            raise ValidationError("arguments_schema must be valid JSON.") from exc
+        return value
+
     def _parse_dependency(self, dep: Any):
         if not isinstance(dep, dict) and not isinstance(dep, str):
             raise ValidationError("'dependencies' should be an array with strings or dict.")
@@ -142,10 +152,7 @@ def _validate_deps(self, deps):
 
     def validate(self, attrs):  # pylint: disable=too-many-branches
         """Validates serializer data."""
-        entrypoint = attrs.get("entrypoint", None)
         image = attrs.get("image", None)
-        if entrypoint is None and image is None:
-            raise ValidationError("At least one of attributes (entrypoint, image) is required.")
         try:
             deps = json.loads(attrs.get("dependencies", "[]"))
         except (json.JSONDecodeError, UnicodeDecodeError) as exc:

gateway/core/models.py

@@ -146,6 +146,7 @@ class Program(ExportModelOperationsMixin("program"), models.Model):
     image = models.CharField(max_length=511, null=True, blank=True)
     env_vars = models.TextField(null=False, blank=True, default="{}")
     dependencies = models.TextField(null=False, blank=True, default="[]")
+    arguments_schema = models.TextField(null=False, blank=True, default="{}")
 
     runner = models.CharField(
         max_length=20, choices=RUNNER_CHOICES, default=RAY, help_text="Execution backend for this program"

gateway/tests/api/test_v1_program.py

@@ -1154,6 +1154,151 @@ def test_upload_without_runner_defaults_to_ray(self):
         program = Program.objects.get(title="Defaultrunnerfunction")
         assert program.runner == Program.RAY
 
+    def test_upload_with_arguments_schema_stores_it(self):
+        """Upload with arguments_schema saves the schema on the Program."""
+        schema = json.dumps({"type": "object", "properties": {"shots": {"type": "integer"}}})
+        fake_file = ContentFile(b"print('hello')")
+        fake_file.name = "test.tar"
+        TestUtils.authorize_client(user="test_user", client=self.client)
+
+        with self.settings(MEDIA_ROOT=self.MEDIA_ROOT):
+            response = self.client.post(
+                "/api/v1/programs/upload/",
+                data={
+                    "title": "schema-function",
+                    "entrypoint": "main.py",
+                    "dependencies": "[]",
+                    "arguments_schema": schema,
+                    "artifact": fake_file,
+                },
+            )
+        assert response.status_code == status.HTTP_200_OK
+        program = Program.objects.get(title="schema-function")
+        assert program.arguments_schema == schema
+
+    def test_upload_with_invalid_schema_returns_400(self):
+        """Upload with malformed JSON as arguments_schema returns 400."""
+        fake_file = ContentFile(b"print('hello')")
+        fake_file.name = "test.tar"
+        TestUtils.authorize_client(user="test_user", client=self.client)
+
+        with self.settings(MEDIA_ROOT=self.MEDIA_ROOT):
+            response = self.client.post(
+                "/api/v1/programs/upload/",
+                data={
+                    "title": "bad-schema-function",
+                    "entrypoint": "main.py",
+                    "dependencies": "[]",
+                    "arguments_schema": "not-valid-json{{{",
+                    "artifact": fake_file,
+                },
+            )
+        assert response.status_code == status.HTTP_400_BAD_REQUEST
+
+    def test_reupload_without_schema_preserves_existing_schema(self):
+        """Re-uploading a function without arguments_schema keeps the existing schema."""
+        schema = json.dumps({"type": "object"})
+        user = TestUtils.authorize_client(user="test_user", client=self.client)
+        TestUtils.create_program(
+            program_title="preserve-schema-func",
+            author=user,
+            entrypoint="main.py",
+            arguments_schema=schema,
+        )
+        fake_file = ContentFile(b"print('hello')")
+        fake_file.name = "test.tar"
+
+        with self.settings(MEDIA_ROOT=self.MEDIA_ROOT):
+            response = self.client.post(
+                "/api/v1/programs/upload/",
+                data={
+                    "title": "preserve-schema-func",
+                    "entrypoint": "main.py",
+                    "dependencies": "[]",
+                    "artifact": fake_file,
+                },
+            )
+        assert response.status_code == status.HTTP_200_OK
+        program = Program.objects.get(title="preserve-schema-func", author=user)
+        assert program.arguments_schema == schema
+
+    def test_upload_all_fields_stores_all_fields(self):
+        """Upload with all optional fields - every field is persisted to the DB."""
+        fake_file = ContentFile(b"print('hello')")
+        fake_file.name = "test.tar"
+        user = TestUtils.authorize_client(user="test_user", client=self.client)
+        schema = json.dumps({"type": "object", "properties": {"n": {"type": "integer"}}})
+        env_vars = json.dumps({"MY_KEY": "MY_VALUE"})
+
+        with self.settings(MEDIA_ROOT=self.MEDIA_ROOT):
+            response = self.client.post(
+                "/api/v1/programs/upload/",
+                data={
+                    "title": "full-function",
+                    "entrypoint": "custom_entry.py",
+                    "dependencies": "[]",
+                    "env_vars": env_vars,
+                    "description": "My description",
+                    "version": "1.2.3",
+                    "runner": Program.RAY,
+                    "arguments_schema": schema,
+                    "artifact": fake_file,
+                },
+            )
+        assert response.status_code == status.HTTP_200_OK
+        program = Program.objects.get(title="full-function", author=user)
+        assert program.entrypoint == "custom_entry.py"
+        assert program.description == "My description"
+        assert program.version == "1.2.3"
+        assert program.runner == Program.RAY
+        assert program.arguments_schema == schema
+        assert program.env_vars not in ("{}", "")
+
+    def test_reupload_with_only_required_fields_preserves_optional_fields(self):
+        """Re-uploading with title+entrypoint must preserve all previously set optional fields.
+
+        This test proves empirically which fields survive a minimal re-upload.
+        Fields that used to RESET (runner, dependencies, env_vars) must be PRESERVED.
+        """
+        schema = json.dumps({"type": "object"})
+        user = TestUtils.authorize_client(user="test_user", client=self.client)
+        TestUtils.get_or_create_ce_project(project_name="test-project", project_id="test-id")
+        TestUtils.create_program(
+            program_title="reupload-test-func",
+            author=user,
+            entrypoint="original.py",
+            description="Original description",
+            version="2.0.0",
+            runner=Program.FLEETS,
+            dependencies='["numpy==1.26.0"]',
+            env_vars='{"MY_KEY":"MY_VALUE"}',
+            arguments_schema=schema,
+        )
+        fake_file = ContentFile(b"print('updated')")
+        fake_file.name = "update.tar"
+
+        with self.settings(MEDIA_ROOT=self.MEDIA_ROOT, CE_DEFAULT_PROJECT_NAME="test-project"):
+            response = self.client.post(
+                "/api/v1/programs/upload/",
+                data={
+                    "title": "reupload-test-func",
+                    "entrypoint": "new_main.py",
+                    "artifact": fake_file,
+                },
+            )
+        assert response.status_code == status.HTTP_200_OK
+        program = Program.objects.get(title="reupload-test-func", author=user)
+
+        # Already preserved (existing behavior):
+        assert program.description == "Original description"
+        assert program.version == "2.0.0"
+        assert program.arguments_schema == schema
+
+        # Now also preserved (previously reset):
+        assert program.runner == Program.FLEETS
+        assert json.loads(program.dependencies) == ["numpy==1.26.0"]
+        assert program.env_vars not in ("{}", "")
+
 
 @pytest.mark.django_db
 class TestProgramApiRuntimeInstances:

gateway/tests/api/test_v1_serializers.py

@@ -257,16 +257,11 @@ def test_run_job_serializer_sets_gpu_flag_for_gpu_provider(self, mock_gpujob_all
         assert job.gpu
 
     def test_upload_program_serializer_with_only_title(self):
-        """Tests upload serializer with only title."""
+        """Serializer accepts title-only requests (entrypoint/image check moved to use case, allowing schema-only updates)."""
         data = {"title": "awesome"}
 
         serializer = UploadProgramSerializer(data=data)
-        assert not serializer.is_valid()
-        errors = serializer.errors
-        assert ["non_field_errors"] == list(errors.keys())
-        assert ["At least one of attributes (entrypoint, image) is required."] == [
-            value[0] for value in errors.values()
-        ]
+        assert serializer.is_valid()
 
     # Dependency validation tests use 'mergedeep' and 'ffsim' as representative examples
     # from requirements-dynamic-dependencies.txt. These tests validate the dependency