cradle.sdk.types.workspace

  1from __future__ import annotations
  2
  3from enum import StrEnum
  4from typing import Literal
  5
  6from pydantic import BaseModel, Field, HttpUrl, RootModel
  7
  8from .common import (
  9    ArchivableResourceMixin,
 10    BaseListResponse,
 11    ErrorResponseMixin,
 12    ImmutableResourceResponse,
 13    ResourceResponse,
 14)
 15
 16
 17class CustomPredictorInputDataType(StrEnum):
 18    BOOL = "BOOL"
 19    INT64 = "INT64"
 20    FLOAT64 = "FLOAT64"
 21    STRING = "STRING"
 22
 23
 24class CustomPredictorKind(StrEnum):
 25    CONTAINER = "CONTAINER"
 26
 27
 28class CustomPredictorState(StrEnum):
 29    ACTIVE = "ACTIVE"
 30    INACTIVE = "INACTIVE"
 31    PROVISIONING = "PROVISIONING"
 32    TESTING = "TESTING"
 33    FAILED = "FAILED"
 34
 35
 36class WorkspaceState(StrEnum):
 37    PROVISIONING = "PROVISIONING"
 38    DELETING = "DELETING"
 39    READY = "READY"
 40    DELETED = "DELETED"
 41    FAILED = "FAILED"
 42
 43
 44class CustomPredictorCreateMixin(BaseModel):
 45    name: str = Field(
 46        description="The custom predictor name with which it is referenced. Must be unique within the workspace and consist of only letters, numbers and the '-' character.",
 47        min_length=1,
 48        max_length=64,
 49    )
 50    display_name: str = Field(description="The human-readable name of the predictor", min_length=1, max_length=255)
 51    description: str = Field(
 52        description="Comprehensive description of what the predictor does", min_length=0, max_length=1024
 53    )
 54    url: HttpUrl | None = Field(
 55        default=None, description="Optional URL where to find more information about the predictor"
 56    )
 57
 58
 59class CustomPredictorUpdateMixin(BaseModel):
 60    inputs: list[CustomPredictorInput] = Field(
 61        description="List of input parameters that the predictor consumes", max_length=64
 62    )
 63    outputs: list[CustomPredictorOutput] = Field(
 64        description="List of outputs that the predictor produces for each sequence", min_length=1, max_length=64
 65    )
 66    batch_size: int = Field(
 67        description="Number of sequences that the predictor should be invoked with. Note that the predictor can be invoked with fewer sequences, but will never be invoked with more."
 68    )
 69    test_configuration: CustomPredictorTestConfiguration | None = Field(
 70        default=None,
 71        description="Configuration for testing the custom predictor. If not `None`, the custom predictor will be invoked with these sequences before becoming available. This makes sure it can run on Cradle's infrastructure and computes the expected outputs.",
 72    )
 73
 74
 75class CustomPredictorContainerMixin(BaseModel):
 76    kind: Literal[CustomPredictorKind.CONTAINER] = Field(default=CustomPredictorKind.CONTAINER)
 77    cpu_mcores: int = Field(description="CPU millicores to allocate for the custom predictor operating on one batch")
 78    main_memory_mib: int = Field(
 79        description="Main memory (RAM) to allocate for the custom predictor operating on one batch"
 80    )
 81    gpu_memory_mib: int = Field(
 82        description="If nonzero, GPU memory to allocate for the custom predictor operating on one batch"
 83    )
 84
 85
 86class CustomPredictorResponseMixin(ImmutableResourceResponse, ErrorResponseMixin):
 87    name: str
 88    display_name: str
 89    description: str
 90    url: str | None
 91    version_id: int
 92    state: CustomPredictorState
 93    inputs: list[CustomPredictorInput]
 94    outputs: list[CustomPredictorOutput]
 95    batch_size: int
 96    test_configuration: CustomPredictorTestConfiguration | None
 97
 98
 99class CustomPredictorCreate(RootModel):
100    root: CustomPredictorContainerCreate = Field(discriminator="kind")
101
102
103class CustomPredictorCreateResponse(RootModel):
104    root: CustomPredictorContainerCreateResponse = Field(discriminator="kind")
105
106
107class CustomPredictorInput(BaseModel):
108    name: str
109    description: str | None = Field(default=None)
110    type: CustomPredictorInputDataType
111    default_value: None | bool | str | int | float = Field(default=None)
112
113
114class CustomPredictorOutput(BaseModel):
115    assay_id: str
116    name: str
117    unit: str | None = Field(default=None)
118
119
120class CustomPredictorResponse(RootModel):
121    root: CustomPredictorContainerResponse = Field(discriminator="kind")
122
123
124class CustomPredictorTestConfiguration(BaseModel):
125    inputs: dict[str, bool | int | float | str] = Field(description="Inputs to use when testing the custom predictor.")
126    sequences: list[CustomPredictorTestSequence] = Field(
127        description="Sequences and expected values to test the custom predictor with. Length must not exceed the predictor's batch size."
128    )
129
130
131class CustomPredictorTestSequence(BaseModel):
132    sequence: str = Field(description="Sequence that the custom predictor will be invoked with.")
133    expected_outputs: dict[str, float | None] = Field(
134        description="Mapping of assay ID to expected value, or `None` if to accept all values and only test that the custom predictor provides a value."
135    )
136    absolute_tolerance: float = Field(
137        default=1e-08, description="Absolute tolerance when comparing the result with the expected output."
138    )
139    relative_tolerance: float = Field(
140        default=1e-05, description="Relative tolerance when comparing the result with the expected output."
141    )
142
143
144class CustomPredictorUpdate(RootModel):
145    root: CustomPredictorContainerUpdate = Field(discriminator="kind")
146
147
148class ListCustomPredictorResponse(BaseListResponse):
149    items: list[CustomPredictorResponse]
150
151
152class ListProjectResponse(BaseListResponse):
153    items: list[ProjectResponse]
154
155
156class ListRoundResponse(BaseListResponse):
157    items: list[RoundResponse]
158
159
160class ProjectCreate(BaseModel):
161    name: str = Field(min_length=1, max_length=64)
162
163
164class ProjectResponse(ResourceResponse, ArchivableResourceMixin):
165    name: str
166
167
168class RoundCreate(BaseModel):
169    project_id: int = Field(description="The project to which the round belongs")
170    name: str = Field(description="The name of the round", min_length=1, max_length=64)
171    description: str | None = Field(default=None, min_length=1, max_length=1024)
172
173
174class RoundResponse(ResourceResponse, ArchivableResourceMixin):
175    project_id: int
176    name: str
177    description: str | None
178
179
180class WorkspaceResponse(ResourceResponse, ErrorResponseMixin):
181    name: str
182    display_name: str
183    state: WorkspaceState
184    picture_file_id: int | None = Field(default=None)
185
186
187class CustomPredictorContainerCreate(
188    CustomPredictorCreateMixin, CustomPredictorUpdateMixin, CustomPredictorContainerMixin
189):
190    pass
191
192
193class CustomPredictorContainerUpdate(CustomPredictorUpdateMixin, CustomPredictorContainerMixin):
194    pass
195
196
197class CustomPredictorContainerCreateResponse(CustomPredictorResponseMixin, CustomPredictorContainerMixin):
198    upload_url: str
199    upload_headers: dict[str, str]
200
201
202class CustomPredictorContainerResponse(CustomPredictorResponseMixin, CustomPredictorContainerMixin):
203    pass