databricks-sdk 0.49.0__py3-none-any.whl → 0.51.0__py3-none-any.whl

databricks/sdk/service/oauth2.py

@@ -1191,10 +1191,10 @@ class AccountFederationPolicyAPI:
     def __init__(self, api_client):
         self._api = api_client
 
-    def create(self, *, policy: Optional[FederationPolicy] = None, policy_id: Optional[str] = None) -> FederationPolicy:
+    def create(self, policy: FederationPolicy, *, policy_id: Optional[str] = None) -> FederationPolicy:
         """Create account federation policy.
 
-        :param policy: :class:`FederationPolicy` (optional)
+        :param policy: :class:`FederationPolicy`
         :param policy_id: str (optional)
           The identifier for the federation policy. The identifier must contain only lowercase alphanumeric
           characters, numbers, hyphens, and slashes. If unspecified, the id will be assigned by Databricks.
@@ -1284,13 +1284,13 @@ class AccountFederationPolicyAPI:
             query["page_token"] = json["next_page_token"]
 
     def update(
-        self, policy_id: str, *, policy: Optional[FederationPolicy] = None, update_mask: Optional[str] = None
+        self, policy_id: str, policy: FederationPolicy, *, update_mask: Optional[str] = None
     ) -> FederationPolicy:
         """Update account federation policy.
 
         :param policy_id: str
           The identifier for the federation policy.
-        :param policy: :class:`FederationPolicy` (optional)
+        :param policy: :class:`FederationPolicy`
         :param update_mask: str (optional)
           The field mask specifies which fields of the policy to update. To specify multiple fields in the
           field mask, use comma as the separator (no space). The special value '*' indicates that all fields
@@ -1758,13 +1758,13 @@ class ServicePrincipalFederationPolicyAPI:
         self._api = api_client
 
     def create(
-        self, service_principal_id: int, *, policy: Optional[FederationPolicy] = None, policy_id: Optional[str] = None
+        self, service_principal_id: int, policy: FederationPolicy, *, policy_id: Optional[str] = None
     ) -> FederationPolicy:
         """Create service principal federation policy.
 
         :param service_principal_id: int
           The service principal id for the federation policy.
-        :param policy: :class:`FederationPolicy` (optional)
+        :param policy: :class:`FederationPolicy`
         :param policy_id: str (optional)
           The identifier for the federation policy. The identifier must contain only lowercase alphanumeric
           characters, numbers, hyphens, and slashes. If unspecified, the id will be assigned by Databricks.
@@ -1869,12 +1869,7 @@ class ServicePrincipalFederationPolicyAPI:
             query["page_token"] = json["next_page_token"]
 
     def update(
-        self,
-        service_principal_id: int,
-        policy_id: str,
-        *,
-        policy: Optional[FederationPolicy] = None,
-        update_mask: Optional[str] = None,
+        self, service_principal_id: int, policy_id: str, policy: FederationPolicy, *, update_mask: Optional[str] = None
     ) -> FederationPolicy:
         """Update service principal federation policy.
 
@@ -1882,7 +1877,7 @@ class ServicePrincipalFederationPolicyAPI:
           The service principal id for the federation policy.
         :param policy_id: str
           The identifier for the federation policy.
-        :param policy: :class:`FederationPolicy` (optional)
+        :param policy: :class:`FederationPolicy`
         :param update_mask: str (optional)
           The field mask specifies which fields of the policy to update. To specify multiple fields in the
           field mask, use comma as the separator (no space). The special value '*' indicates that all fields

databricks/sdk/service/pipelines.py

@@ -3390,38 +3390,6 @@ class PipelinesAPI:
     def __init__(self, api_client):
         self._api = api_client
 
-    def wait_get_pipeline_running(
-        self,
-        pipeline_id: str,
-        timeout=timedelta(minutes=20),
-        callback: Optional[Callable[[GetPipelineResponse], None]] = None,
-    ) -> GetPipelineResponse:
-        deadline = time.time() + timeout.total_seconds()
-        target_states = (PipelineState.RUNNING,)
-        failure_states = (PipelineState.FAILED,)
-        status_message = "polling..."
-        attempt = 1
-        while time.time() < deadline:
-            poll = self.get(pipeline_id=pipeline_id)
-            status = poll.state
-            status_message = poll.cause
-            if status in target_states:
-                return poll
-            if callback:
-                callback(poll)
-            if status in failure_states:
-                msg = f"failed to reach RUNNING, got {status}: {status_message}"
-                raise OperationFailed(msg)
-            prefix = f"pipeline_id={pipeline_id}"
-            sleep = attempt
-            if sleep > 10:
-                # sleep 10s max per attempt
-                sleep = 10
-            _LOG.debug(f"{prefix}: ({status}) {status_message} (sleeping ~{sleep}s)")
-            time.sleep(sleep + random.random())
-            attempt += 1
-        raise TimeoutError(f"timed out after {timeout}: {status_message}")
-
     def wait_get_pipeline_idle(
         self,
         pipeline_id: str,

databricks/sdk/service/serving.py

@@ -849,8 +849,9 @@ class CreateServingEndpoint:
     workspace. An endpoint name can consist of alphanumeric characters, dashes, and underscores."""
 
     ai_gateway: Optional[AiGatewayConfig] = None
-    """The AI Gateway configuration for the serving endpoint. NOTE: Only external model and provisioned
-    throughput endpoints are currently supported."""
+    """The AI Gateway configuration for the serving endpoint. NOTE: External model, provisioned
+    throughput, and pay-per-token endpoints are fully supported; agent endpoints currently only
+    support inference tables."""
 
     budget_policy_id: Optional[str] = None
     """The budget policy to be applied to the serving endpoint."""
@@ -2873,7 +2874,8 @@ class ServedEntityInput:
     """The workload size of the served entity. The workload size corresponds to a range of provisioned
     concurrency that the compute autoscales between. A single unit of provisioned concurrency can
     process one request at a time. Valid workload sizes are "Small" (4 - 4 provisioned concurrency),
-    "Medium" (8 - 16 provisioned concurrency), and "Large" (16 - 64 provisioned concurrency). If
+    "Medium" (8 - 16 provisioned concurrency), and "Large" (16 - 64 provisioned concurrency).
+    Additional custom workload sizes can also be used when available in the workspace. If
     scale-to-zero is enabled, the lower bound of the provisioned concurrency for each workload size
     is 0."""
 
@@ -3013,7 +3015,8 @@ class ServedEntityOutput:
     """The workload size of the served entity. The workload size corresponds to a range of provisioned
     concurrency that the compute autoscales between. A single unit of provisioned concurrency can
     process one request at a time. Valid workload sizes are "Small" (4 - 4 provisioned concurrency),
-    "Medium" (8 - 16 provisioned concurrency), and "Large" (16 - 64 provisioned concurrency). If
+    "Medium" (8 - 16 provisioned concurrency), and "Large" (16 - 64 provisioned concurrency).
+    Additional custom workload sizes can also be used when available in the workspace. If
     scale-to-zero is enabled, the lower bound of the provisioned concurrency for each workload size
     is 0."""
 
@@ -3203,11 +3206,12 @@ class ServedModelInput:
     model, this field defaults to external_model.name, with '.' and ':' replaced with '-', and if
     not specified for other entities, it defaults to entity_name-entity_version."""
 
-    workload_size: Optional[ServedModelInputWorkloadSize] = None
+    workload_size: Optional[str] = None
     """The workload size of the served entity. The workload size corresponds to a range of provisioned
     concurrency that the compute autoscales between. A single unit of provisioned concurrency can
     process one request at a time. Valid workload sizes are "Small" (4 - 4 provisioned concurrency),
-    "Medium" (8 - 16 provisioned concurrency), and "Large" (16 - 64 provisioned concurrency). If
+    "Medium" (8 - 16 provisioned concurrency), and "Large" (16 - 64 provisioned concurrency).
+    Additional custom workload sizes can also be used when available in the workspace. If
     scale-to-zero is enabled, the lower bound of the provisioned concurrency for each workload size
     is 0."""
 
@@ -3239,7 +3243,7 @@ class ServedModelInput:
         if self.scale_to_zero_enabled is not None:
             body["scale_to_zero_enabled"] = self.scale_to_zero_enabled
         if self.workload_size is not None:
-            body["workload_size"] = self.workload_size.value
+            body["workload_size"] = self.workload_size
         if self.workload_type is not None:
             body["workload_type"] = self.workload_type.value
         return body
@@ -3281,18 +3285,11 @@ class ServedModelInput:
             model_version=d.get("model_version", None),
             name=d.get("name", None),
             scale_to_zero_enabled=d.get("scale_to_zero_enabled", None),
-            workload_size=_enum(d, "workload_size", ServedModelInputWorkloadSize),
+            workload_size=d.get("workload_size", None),
             workload_type=_enum(d, "workload_type", ServedModelInputWorkloadType),
         )
 
 
-class ServedModelInputWorkloadSize(Enum):
-
-    LARGE = "Large"
-    MEDIUM = "Medium"
-    SMALL = "Small"
-
-
 class ServedModelInputWorkloadType(Enum):
     """Please keep this in sync with with workload types in InferenceEndpointEntities.scala"""
 
@@ -3337,7 +3334,8 @@ class ServedModelOutput:
     """The workload size of the served entity. The workload size corresponds to a range of provisioned
     concurrency that the compute autoscales between. A single unit of provisioned concurrency can
     process one request at a time. Valid workload sizes are "Small" (4 - 4 provisioned concurrency),
-    "Medium" (8 - 16 provisioned concurrency), and "Large" (16 - 64 provisioned concurrency). If
+    "Medium" (8 - 16 provisioned concurrency), and "Large" (16 - 64 provisioned concurrency).
+    Additional custom workload sizes can also be used when available in the workspace. If
     scale-to-zero is enabled, the lower bound of the provisioned concurrency for each workload size
     is 0."""
 
@@ -3531,8 +3529,9 @@ class ServerLogsResponse:
 @dataclass
 class ServingEndpoint:
     ai_gateway: Optional[AiGatewayConfig] = None
-    """The AI Gateway configuration for the serving endpoint. NOTE: Only external model and provisioned
-    throughput endpoints are currently supported."""
+    """The AI Gateway configuration for the serving endpoint. NOTE: External model, provisioned
+    throughput, and pay-per-token endpoints are fully supported; agent endpoints currently only
+    support inference tables."""
 
     budget_policy_id: Optional[str] = None
     """The budget policy associated with the endpoint."""
@@ -3749,8 +3748,9 @@ class ServingEndpointAccessControlResponse:
 @dataclass
 class ServingEndpointDetailed:
     ai_gateway: Optional[AiGatewayConfig] = None
-    """The AI Gateway configuration for the serving endpoint. NOTE: Only external model and provisioned
-    throughput endpoints are currently supported."""
+    """The AI Gateway configuration for the serving endpoint. NOTE: External model, provisioned
+    throughput, and pay-per-token endpoints are fully supported; agent endpoints currently only
+    support inference tables."""
 
     budget_policy_id: Optional[str] = None
     """The budget policy associated with the endpoint."""
@@ -4242,8 +4242,9 @@ class ServingEndpointsAPI:
           The name of the serving endpoint. This field is required and must be unique across a Databricks
           workspace. An endpoint name can consist of alphanumeric characters, dashes, and underscores.
         :param ai_gateway: :class:`AiGatewayConfig` (optional)
-          The AI Gateway configuration for the serving endpoint. NOTE: Only external model and provisioned
-          throughput endpoints are currently supported.
+          The AI Gateway configuration for the serving endpoint. NOTE: External model, provisioned throughput,
+          and pay-per-token endpoints are fully supported; agent endpoints currently only support inference
+          tables.
         :param budget_policy_id: str (optional)
           The budget policy to be applied to the serving endpoint.
         :param config: :class:`EndpointCoreConfigInput` (optional)
@@ -4533,8 +4534,7 @@ class ServingEndpointsAPI:
     def put(self, name: str, *, rate_limits: Optional[List[RateLimit]] = None) -> PutResponse:
         """Update rate limits of a serving endpoint.
 
-        Used to update the rate limits of a serving endpoint. NOTE: Only foundation model endpoints are
-        currently supported. For external models, use AI Gateway to manage rate limits.
+        Deprecated: Please use AI Gateway to manage rate limits instead.
 
         :param name: str
           The name of the serving endpoint whose rate limits are being updated. This field is required.
@@ -4566,8 +4566,8 @@ class ServingEndpointsAPI:
     ) -> PutAiGatewayResponse:
         """Update AI Gateway of a serving endpoint.
 
-        Used to update the AI Gateway of a serving endpoint. NOTE: Only external model and provisioned
-        throughput endpoints are currently supported.
+        Used to update the AI Gateway of a serving endpoint. NOTE: External model, provisioned throughput, and
+        pay-per-token endpoints are fully supported; agent endpoints currently only support inference tables.
 
         :param name: str
           The name of the serving endpoint whose AI Gateway is being updated. This field is required.