From 4ca8a4e0454c28d05d63f264cd7b22ae7430ab7b Mon Sep 17 00:00:00 2001 From: oatanas Date: Fri, 11 Sep 2026 13:34:41 +0200 Subject: [PATCH] feat(platform): add RTX_PRO_6000 to the GPUType enum MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The he-tme application accepts RTX PRO 6000 (sm_120) GPUs, but the SDK enum stopped at L4 and A100, so `validate_run_sdk_metadata` — which `Runs.submit` calls on every submission — rejected the run before it reached the platform: pipeline.gpu.gpu_type Input should be 'L4' or 'A100' [type=enum, input_value='RTX_PRO_6000'] That gate is inside the SDK and `submit` has no opt-out, so no caller could work around it; passing the string through from a downstream CLI still failed here. Bumps SDK_METADATA_SCHEMA_VERSION to 0.0.7 and regenerates the published run schema, following 0.0.6, which was bumped for the equivalent change (adding FLEX_START to the provisioning-mode enum). Adding an enum member widens what validates, so runs written against 0.0.6 stay valid. CLI_REFERENCE.md is left alone: it is generated, has drifted from the current CLI, and `nox -s docs` regenerates it with all extras installed. Refs HETA-1792 Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01CNyXyfHNFAgv7zUYjUG22P --- ...sdk_run_custom_metadata_schema_latest.json | 7 +- ...sdk_run_custom_metadata_schema_v0.0.7.json | 587 ++++++++++++++++++ src/aignostics_sdk/application/_cli.py | 2 +- src/aignostics_sdk/application/_service.py | 4 +- src/aignostics_sdk/platform/_sdk_metadata.py | 5 +- .../aignostics/platform/sdk_metadata_test.py | 3 +- 6 files changed, 599 insertions(+), 9 deletions(-) create mode 100644 docs/source/_static/sdk_run_custom_metadata_schema_v0.0.7.json diff --git a/docs/source/_static/sdk_run_custom_metadata_schema_latest.json b/docs/source/_static/sdk_run_custom_metadata_schema_latest.json index 8092381c2..205add658 100644 --- a/docs/source/_static/sdk_run_custom_metadata_schema_latest.json +++ b/docs/source/_static/sdk_run_custom_metadata_schema_latest.json @@ -47,7 +47,7 @@ "properties": { "gpu_type": { "$ref": "#/$defs/GPUType", - "description": "The type of GPU to use (L4 or A100)" + "description": "The type of GPU to use (L4, A100 or RTX_PRO_6000)" }, "provisioning_mode": { "$ref": "#/$defs/ProvisioningMode", @@ -84,7 +84,8 @@ "description": "Type of GPU to use for processing.", "enum": [ "L4", - "A100" + "A100", + "RTX_PRO_6000" ], "title": "GPUType", "type": "string" @@ -582,5 +583,5 @@ "title": "RunSdkMetadata", "type": "object", "$schema": "https://json-schema.org/draft/2020-12/schema", - "$id": "https://raw.githubusercontent.com/aignostics/python-sdk/main/docs/source/_static/sdk_metadata_schema_v0.0.6.json" + "$id": "https://raw.githubusercontent.com/aignostics/python-sdk/main/docs/source/_static/sdk_metadata_schema_v0.0.7.json" } \ No newline at end of file diff --git a/docs/source/_static/sdk_run_custom_metadata_schema_v0.0.7.json b/docs/source/_static/sdk_run_custom_metadata_schema_v0.0.7.json new file mode 100644 index 000000000..205add658 --- /dev/null +++ b/docs/source/_static/sdk_run_custom_metadata_schema_v0.0.7.json @@ -0,0 +1,587 @@ +{ + "$defs": { + "CIMetadata": { + "description": "CI/CD environment metadata.", + "properties": { + "github": { + "anyOf": [ + { + "$ref": "#/$defs/GitHubCIMetadata" + }, + { + "type": "null" + } + ], + "default": null, + "description": "GitHub Actions metadata" + }, + "pytest": { + "anyOf": [ + { + "$ref": "#/$defs/PytestCIMetadata" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Pytest test metadata" + } + }, + "title": "CIMetadata", + "type": "object" + }, + "CPUConfig": { + "description": "Configuration for CPU resources.", + "properties": { + "provisioning_mode": { + "$ref": "#/$defs/ProvisioningMode", + "description": "The provisioning mode for CPU resources (SPOT or ON_DEMAND)" + } + }, + "title": "CPUConfig", + "type": "object" + }, + "GPUConfig": { + "description": "Configuration for GPU resources.", + "properties": { + "gpu_type": { + "$ref": "#/$defs/GPUType", + "description": "The type of GPU to use (L4, A100 or RTX_PRO_6000)" + }, + "provisioning_mode": { + "$ref": "#/$defs/ProvisioningMode", + "description": "The provisioning mode for GPU resources (SPOT, ON_DEMAND, or FLEX_START)" + }, + "max_gpus_per_slide": { + "default": 1, + "description": "The maximum number of GPUs to allocate per slide (1-8)", + "maximum": 8, + "minimum": 1, + "title": "Max Gpus Per Slide", + "type": "integer" + }, + "flex_start_max_run_duration_minutes": { + "anyOf": [ + { + "maximum": 3600, + "minimum": 1, + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Maximum run duration in minutes when using FLEX_START provisioning mode (1-3600). Required when provisioning_mode is FLEX_START, must be None otherwise.", + "title": "Flex Start Max Run Duration Minutes" + } + }, + "title": "GPUConfig", + "type": "object" + }, + "GPUType": { + "description": "Type of GPU to use for processing.", + "enum": [ + "L4", + "A100", + "RTX_PRO_6000" + ], + "title": "GPUType", + "type": "string" + }, + "GitHubCIMetadata": { + "description": "GitHub Actions CI metadata.", + "properties": { + "action": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "GitHub Action name", + "title": "Action" + }, + "job": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "GitHub job name", + "title": "Job" + }, + "ref": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Git reference", + "title": "Ref" + }, + "ref_name": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Git reference name", + "title": "Ref Name" + }, + "ref_type": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Git reference type (branch, tag)", + "title": "Ref Type" + }, + "repository": { + "description": "Repository name (owner/repo)", + "title": "Repository", + "type": "string" + }, + "run_attempt": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Attempt number for this run", + "title": "Run Attempt" + }, + "run_id": { + "description": "Unique ID for this workflow run", + "title": "Run Id", + "type": "string" + }, + "run_number": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Run number for this workflow", + "title": "Run Number" + }, + "run_url": { + "description": "URL to the workflow run", + "title": "Run Url", + "type": "string" + }, + "runner_arch": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Runner architecture (x64, ARM64, etc.)", + "title": "Runner Arch" + }, + "runner_os": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Runner operating system", + "title": "Runner Os" + }, + "sha": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Git commit SHA", + "title": "Sha" + }, + "workflow": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Workflow name", + "title": "Workflow" + }, + "workflow_ref": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Reference to the workflow file", + "title": "Workflow Ref" + } + }, + "required": [ + "repository", + "run_id", + "run_url" + ], + "title": "GitHubCIMetadata", + "type": "object" + }, + "PipelineConfig": { + "description": "Pipeline configuration for dynamic orchestration.", + "properties": { + "gpu": { + "$ref": "#/$defs/GPUConfig", + "description": "GPU resource configuration" + }, + "cpu": { + "$ref": "#/$defs/CPUConfig", + "description": "CPU resource configuration" + }, + "node_acquisition_timeout_minutes": { + "default": 30, + "description": "Timeout for acquiring compute nodes in minutes (1-3600)", + "maximum": 3600, + "minimum": 1, + "title": "Node Acquisition Timeout Minutes", + "type": "integer" + } + }, + "title": "PipelineConfig", + "type": "object" + }, + "ProvisioningMode": { + "description": "Provisioning mode for resources.", + "enum": [ + "SPOT", + "ON_DEMAND", + "FLEX_START" + ], + "title": "ProvisioningMode", + "type": "string" + }, + "PytestCIMetadata": { + "description": "Pytest test execution metadata.", + "properties": { + "current_test": { + "description": "Current test being executed", + "title": "Current Test", + "type": "string" + }, + "markers": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Pytest markers applied to the test", + "title": "Markers" + } + }, + "required": [ + "current_test" + ], + "title": "PytestCIMetadata", + "type": "object" + }, + "SchedulingMetadata": { + "description": "Scheduling metadata for run execution.", + "properties": { + "due_date": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Requested completion time (ISO 8601). Scheduler will try to complete before this time.", + "title": "Due Date" + }, + "deadline": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Hard deadline (ISO 8601). Run may be aborted if processing exceeds this time.", + "title": "Deadline" + } + }, + "title": "SchedulingMetadata", + "type": "object" + }, + "SubmissionMetadata": { + "description": "Metadata about how the SDK was invoked.", + "properties": { + "date": { + "description": "ISO 8601 timestamp of submission", + "title": "Date", + "type": "string" + }, + "interface": { + "description": "How the SDK was accessed (script, cli, launchpad)", + "enum": [ + "script", + "cli", + "launchpad" + ], + "title": "Interface", + "type": "string" + }, + "initiator": { + "description": "Who/what initiated the run (user, test, bridge)", + "enum": [ + "user", + "test", + "bridge" + ], + "title": "Initiator", + "type": "string" + } + }, + "required": [ + "date", + "interface", + "initiator" + ], + "title": "SubmissionMetadata", + "type": "object" + }, + "UserMetadata": { + "description": "User information metadata.", + "properties": { + "organization_id": { + "description": "User's organization ID", + "title": "Organization Id", + "type": "string" + }, + "organization_name": { + "description": "User's organization name", + "title": "Organization Name", + "type": "string" + }, + "user_email": { + "description": "User's email address", + "title": "User Email", + "type": "string" + }, + "user_id": { + "description": "User's unique ID", + "title": "User Id", + "type": "string" + } + }, + "required": [ + "organization_id", + "organization_name", + "user_email", + "user_id" + ], + "title": "UserMetadata", + "type": "object" + }, + "WorkflowMetadata": { + "description": "Workflow control metadata.", + "properties": { + "onboard_to_aignostics_portal": { + "default": false, + "description": "Whether to onboard results to the Aignostics Portal", + "title": "Onboard To Aignostics Portal", + "type": "boolean" + } + }, + "title": "WorkflowMetadata", + "type": "object" + } + }, + "additionalProperties": false, + "description": "Complete Run SDK metadata schema.\n\nThis model defines the structure and validation rules for SDK metadata\nthat is attached to application runs. It includes information about:\n- SDK version and timestamps\n- User information (when available)\n- CI/CD environment context (GitHub Actions, pytest)\n- Workflow control flags\n- Scheduling information\n- Optional user note", + "properties": { + "schema_version": { + "description": "Schema version for this metadata format", + "pattern": "^\\d+\\.\\d+\\.\\d+-?.*$", + "title": "Schema Version", + "type": "string" + }, + "created_at": { + "description": "ISO 8601 timestamp when the metadata was first created", + "title": "Created At", + "type": "string" + }, + "updated_at": { + "description": "ISO 8601 timestamp when the metadata was last updated", + "title": "Updated At", + "type": "string" + }, + "tags": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": true + }, + { + "type": "null" + } + ], + "default": null, + "description": "Optional list of tags associated with the run", + "title": "Tags" + }, + "submission": { + "$ref": "#/$defs/SubmissionMetadata", + "description": "Submission context metadata" + }, + "user_agent": { + "description": "User agent string for the SDK client", + "title": "User Agent", + "type": "string" + }, + "user": { + "anyOf": [ + { + "$ref": "#/$defs/UserMetadata" + }, + { + "type": "null" + } + ], + "default": null, + "description": "User information (when authenticated)" + }, + "ci": { + "anyOf": [ + { + "$ref": "#/$defs/CIMetadata" + }, + { + "type": "null" + } + ], + "default": null, + "description": "CI/CD environment metadata" + }, + "note": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Optional user note for the run", + "title": "Note" + }, + "workflow": { + "anyOf": [ + { + "$ref": "#/$defs/WorkflowMetadata" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Workflow control flags" + }, + "scheduling": { + "anyOf": [ + { + "$ref": "#/$defs/SchedulingMetadata" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Scheduling information" + }, + "pipeline": { + "anyOf": [ + { + "$ref": "#/$defs/PipelineConfig" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Pipeline orchestration configuration" + } + }, + "required": [ + "schema_version", + "created_at", + "updated_at", + "submission", + "user_agent" + ], + "title": "RunSdkMetadata", + "type": "object", + "$schema": "https://json-schema.org/draft/2020-12/schema", + "$id": "https://raw.githubusercontent.com/aignostics/python-sdk/main/docs/source/_static/sdk_metadata_schema_v0.0.7.json" +} \ No newline at end of file diff --git a/src/aignostics_sdk/application/_cli.py b/src/aignostics_sdk/application/_cli.py index 4a0d355ac..2eafb49e2 100644 --- a/src/aignostics_sdk/application/_cli.py +++ b/src/aignostics_sdk/application/_cli.py @@ -102,7 +102,7 @@ def _Service() -> Service: # noqa: N802 GpuTypeOption = Annotated[ str, - typer.Option(help="GPU type to use for processing (L4 or A100)."), + typer.Option(help="GPU type to use for processing (L4, A100 or RTX_PRO_6000)."), ] GpuProvisioningModeOption = Annotated[ diff --git a/src/aignostics_sdk/application/_service.py b/src/aignostics_sdk/application/_service.py index 48be4cf55..a44e69b87 100644 --- a/src/aignostics_sdk/application/_service.py +++ b/src/aignostics_sdk/application/_service.py @@ -854,7 +854,7 @@ def application_run_submit_from_metadata( # noqa: PLR0913 application_version (str | None): The version of the application. If not given latest version is used. onboard_to_aignostics_portal (bool): True if the run should be onboarded to the Aignostics Portal. - gpu_type (str | None): The type of GPU to use (L4 or A100). + gpu_type (str | None): The type of GPU to use (L4, A100 or RTX_PRO_6000). gpu_provisioning_mode (str | None): The provisioning mode for GPU resources (SPOT, ON_DEMAND, or FLEX_START). max_gpus_per_slide (int | None): The maximum number of GPUs to allocate per slide. @@ -1023,7 +1023,7 @@ def application_run_submit( # noqa: PLR0913, PLR0912, C901, PLR0915 deadline (str | None): An optional hard deadline for the run, ISO8601 format. If processing exceeds this deadline, the run can be aborted. onboard_to_aignostics_portal (bool): True if the run should be onboarded to the Aignostics Portal. - gpu_type (str | None): The type of GPU to use (L4 or A100). + gpu_type (str | None): The type of GPU to use (L4, A100 or RTX_PRO_6000). gpu_provisioning_mode (str | None): The provisioning mode for GPU resources (SPOT, ON_DEMAND, or FLEX_START). max_gpus_per_slide (int | None): The maximum number of GPUs to allocate per slide. diff --git a/src/aignostics_sdk/platform/_sdk_metadata.py b/src/aignostics_sdk/platform/_sdk_metadata.py index 80f25d1e7..de23d8fcd 100644 --- a/src/aignostics_sdk/platform/_sdk_metadata.py +++ b/src/aignostics_sdk/platform/_sdk_metadata.py @@ -24,7 +24,7 @@ DEFAULT_NODE_ACQUISITION_TIMEOUT_MINUTES, ) -SDK_METADATA_SCHEMA_VERSION = "0.0.6" +SDK_METADATA_SCHEMA_VERSION = "0.0.7" ITEM_SDK_METADATA_SCHEMA_VERSION = "0.0.3" VALIDATION_CASE_TAG_PREFIX = "__aignx_validation_case:" @@ -34,6 +34,7 @@ class GPUType(StrEnum): L4 = "L4" A100 = "A100" + RTX_PRO_6000 = "RTX_PRO_6000" class ProvisioningMode(StrEnum): @@ -68,7 +69,7 @@ class GPUConfig(BaseModel): gpu_type: GPUType = Field( default_factory=lambda: GPUType(DEFAULT_GPU_TYPE), - description="The type of GPU to use (L4 or A100)", + description="The type of GPU to use (L4, A100 or RTX_PRO_6000)", ) provisioning_mode: ProvisioningMode = Field( default_factory=lambda: ProvisioningMode(DEFAULT_GPU_PROVISIONING_MODE), diff --git a/tests/aignostics/platform/sdk_metadata_test.py b/tests/aignostics/platform/sdk_metadata_test.py index 4de089c1c..78f538351 100644 --- a/tests/aignostics/platform/sdk_metadata_test.py +++ b/tests/aignostics/platform/sdk_metadata_test.py @@ -1037,7 +1037,8 @@ def test_gpu_type_enum() -> None: assert GPUType.L4.value == "L4" assert GPUType.A100.value == "A100" - assert len(GPUType) == 2 + assert GPUType.RTX_PRO_6000.value == "RTX_PRO_6000" + assert len(GPUType) == 3 @pytest.mark.unit @staticmethod