diff --git a/.codegen/_last_sha b/.codegen/_last_sha index 6f2f696f1..d80d0c353 100644 --- a/.codegen/_last_sha +++ b/.codegen/_last_sha @@ -1 +1 @@ -231772f37bb5c01154bfc3d1401dc241ff1b7409 \ No newline at end of file +bcd995071117be59fe1e8c3e3ecfae77ac90b355 \ No newline at end of file diff --git a/.github/workflows/push.yml b/.github/workflows/push.yml index a80d39885..63fb5c7a9 100644 --- a/.github/workflows/push.yml +++ b/.github/workflows/push.yml @@ -31,35 +31,6 @@ jobs: os: windows-latest pyVersion: ${{ matrix.pyVersion }} - fmt: - runs-on: ubuntu-latest - - steps: - - name: Checkout - uses: actions/checkout@ee0669bd1cc54295c223e0bb666b733df41de1c5 # v2.7.0 - - - name: Install uv - uses: astral-sh/setup-uv@eac588ad8def6316056a12d4907a9d4d84ff7a3b # v7.3.0 - with: - version: "0.6.5" - - - name: Setup JFrog PyPI proxy - uses: ./.github/actions/setup-jfrog-pypi - - - name: Re-lock for JFrog - run: uv lock - - - name: Format all files - run: make dev fmt - - - name: Restore lockfiles and fix proxy URLs - run: | - git checkout -- '*.lock' - make fix-lockfile - - - name: Fail on differences - run: git diff --exit-code - check-manifest: runs-on: ubuntu-latest diff --git a/CHANGELOG.md b/CHANGELOG.md index 11f5283f5..90a9f2413 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,22 @@ # Version changelog +## Release v0.144.0 (2026-09-30) + +### API Changes +* Add `ai_decide()` method for [w.ai_functions](https://databricks-sdk-py.readthedocs.io/en/latest/workspace/aifunctions/ai_functions.html) workspace-level service. +* Add `environment_variables` field for `databricks.sdk.service.jobs.BaseRun`. +* Add `environment_variables` field for `databricks.sdk.service.jobs.CreateJob`. +* Add `environment_variables` field for `databricks.sdk.service.jobs.JobSettings`. +* Add `environment_variables` field for `databricks.sdk.service.jobs.Run`. +* Add `environment_variables_key` field for `databricks.sdk.service.jobs.RunTask`. +* Add `environment_variables` field for `databricks.sdk.service.jobs.SubmitRun`. +* Add `environment_variables_key` field for `databricks.sdk.service.jobs.SubmitTask`. +* Add `environment_variables_key` field for `databricks.sdk.service.jobs.Task`. +* [Breaking] Change `command_path` field for `databricks.sdk.service.jobs.DeploymentSpec` to no longer be required. +* Change `command_path` field for `databricks.sdk.service.jobs.DeploymentSpec` to no longer be required. +* Change `api_secret_ref` field for `databricks.sdk.service.ml.SchemaRegistryConfig` to no longer be required. +* [Breaking] Change `api_secret_ref` field for `databricks.sdk.service.ml.SchemaRegistryConfig` to no longer be required. + ## Release v0.143.0 (2026-09-26) ### API Changes diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 2b604d5aa..cca9d3257 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -31,7 +31,7 @@ Small patches and bug fixes don't need prior communication. ## Coding Style -Code style is enforced by a formatter check in your pull request. We use [Black](https://github.com/psf/black) to format our code. Run `make fmt` to ensure your code is properly formatted prior to raising a pull request. +We use Ruff for Python formatting and linting. Run `make fmt` to format code and apply automatic lint fixes. Run `make lint` to check for remaining lint and formatting issues before raising a pull request. ## Signed Commits This repo requires all contributors to sign their commits. To configure this, you can follow [Github's documentation](https://docs.github.com/en/authentication/managing-commit-signature-verification/signing-commits) to create a GPG key, upload it to your Github account, and configure your git client to sign commits. diff --git a/databricks/sdk/__init__.py b/databricks/sdk/__init__.py index 5e670ff81..ea3367399 100644 --- a/databricks/sdk/__init__.py +++ b/databricks/sdk/__init__.py @@ -55,6 +55,7 @@ from databricks.sdk.mixins.workspace import WorkspaceExt from databricks.sdk.mixins.open_ai_client import ServingEndpointsExt from databricks.sdk.mixins.jobs import JobsExt +from databricks.sdk.mixins.sandbox import SandboxExt from databricks.sdk.oauth import AuthorizationDetail from databricks.sdk.service.iam import AccessControlAPI from databricks.sdk.service.iam import AccountAccessControlAPI @@ -468,7 +469,7 @@ def __init__( self._repos = pkg_workspace.ReposAPI(self._api_client) self._resource_quotas = pkg_catalog.ResourceQuotasAPI(self._api_client) self._rfa = pkg_catalog.RfaAPI(self._api_client) - self._sandbox = pkg_sandbox.SandboxAPI(self._api_client) + self._sandbox = SandboxExt(self._api_client) self._schemas = pkg_catalog.SchemasAPI(self._api_client) self._secrets = pkg_workspace.SecretsAPI(self._api_client) self._secrets_uc = pkg_catalog.SecretsUcAPI(self._api_client) @@ -1031,7 +1032,7 @@ def rfa(self) -> pkg_catalog.RfaAPI: return self._rfa @property - def sandbox(self) -> pkg_sandbox.SandboxAPI: + def sandbox(self) -> SandboxExt: """Create, manage, and control the lifecycle of sandboxes -- isolated, pre-configured, low-latency Serverless compute environments for running code.""" return self._sandbox diff --git a/databricks/sdk/_sandbox_grpc/__init__.py b/databricks/sdk/_sandbox_grpc/__init__.py new file mode 100644 index 000000000..e75a6d8ba --- /dev/null +++ b/databricks/sdk/_sandbox_grpc/__init__.py @@ -0,0 +1,28 @@ +"""Committed gRPC stubs for the Sandbox command-execution service. + +sandbox_service_pb2 / sandbox_service_pb2_grpc are generated from the shared +client proto (kept out of this package) at +compute-fabric/sandbox-sdk/proto/databricks/sandbox/sandbox_service.proto, a +dependency-stripped copy of +compute-fabric/sandbox-daemon/proto/sandbox_service.proto. The mixin uses only the +streaming RPCs (ExecuteCommand, AttachCommand, StreamInput); the stub also carries +the unary ListCommands, which the mixin does not use. Stubs are committed, not +built: the daemon proto pulls in internal imports that cannot ship publicly. + +SDK codegen can't emit streaming clients yet. TODO(XTA-19715): migrate the mixin to +native streaming codegen and drop this package. + +Regenerate after the shared proto changes, with a grpcio-tools whose bundled protoc +matches the pinned protobuf major (the sandbox extra in pyproject.toml). Stage the +proto at this package's path so the module path stays databricks.sdk._sandbox_grpc: + + root=$(mktemp -d) + mkdir -p "$root/databricks/sdk/_sandbox_grpc" + cp /compute-fabric/sandbox-sdk/proto/databricks/sandbox/sandbox_service.proto \\ + "$root/databricks/sdk/_sandbox_grpc/sandbox_service.proto" + python -m grpc_tools.protoc -I "$root" \\ + --python_out="$root" --grpc_python_out="$root" \\ + databricks/sdk/_sandbox_grpc/sandbox_service.proto + cp "$root"/databricks/sdk/_sandbox_grpc/sandbox_service_pb2*.py \\ + databricks/sdk/_sandbox_grpc/ +""" diff --git a/databricks/sdk/_sandbox_grpc/sandbox_service_pb2.py b/databricks/sdk/_sandbox_grpc/sandbox_service_pb2.py new file mode 100644 index 000000000..6d39fcfd4 --- /dev/null +++ b/databricks/sdk/_sandbox_grpc/sandbox_service_pb2.py @@ -0,0 +1,74 @@ +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# NO CHECKED-IN PROTOBUF GENCODE +# source: databricks/sdk/_sandbox_grpc/sandbox_service.proto +# Protobuf Python Version: 6.33.5 +"""Generated protocol buffer code.""" +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import runtime_version as _runtime_version +from google.protobuf import symbol_database as _symbol_database +from google.protobuf.internal import builder as _builder +_runtime_version.ValidateProtobufRuntimeVersion( + _runtime_version.Domain.PUBLIC, + 6, + 33, + 5, + '', + 'databricks/sdk/_sandbox_grpc/sandbox_service.proto' +) +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n2databricks/sdk/_sandbox_grpc/sandbox_service.proto\x12\x12\x64\x61tabricks.sandbox\"\xe2\x01\n\x15\x45xecuteCommandRequest\x12\x0b\n\x03\x63md\x18\x01 \x01(\t\x12\x0c\n\x04\x61rgs\x18\x02 \x03(\t\x12*\n\x03pty\x18\x03 \x01(\x0b\x32\x1d.databricks.sandbox.PtyConfig\x12\x41\n\x04\x65nvs\x18\x04 \x03(\x0b\x32\x33.databricks.sandbox.ExecuteCommandRequest.EnvsEntry\x12\x12\n\ncommand_id\x18\x05 \x01(\t\x1a+\n\tEnvsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\"H\n\x14\x41ttachCommandRequest\x12\x12\n\ncommand_id\x18\x01 \x01(\t\x12\x1c\n\x14last_sequence_number\x18\x02 \x01(\x03\"\xc0\x01\n\x12StreamInputRequest\x12\x35\n\x05start\x18\x01 \x01(\x0b\x32$.databricks.sandbox.StreamInputStartH\x00\x12\x33\n\x04\x64\x61ta\x18\x02 \x01(\x0b\x32#.databricks.sandbox.StreamInputDataH\x00\x12\x35\n\x05\x63lose\x18\x03 \x01(\x0b\x32$.databricks.sandbox.StreamInputCloseH\x00\x42\x07\n\x05\x65vent\"&\n\x10StreamInputStart\x12\x12\n\ncommand_id\x18\x01 \x01(\t\"\x1f\n\x0fStreamInputData\x12\x0c\n\x04\x64\x61ta\x18\x01 \x01(\x0c\"\x12\n\x10StreamInputClose\"\x15\n\x13StreamInputResponse\"<\n\x13ListCommandsRequest\x12\x11\n\tpage_size\x18\x01 \x01(\x05\x12\x12\n\npage_token\x18\x02 \x01(\t\"^\n\x14ListCommandsResponse\x12-\n\x08\x63ommands\x18\x01 \x03(\x0b\x32\x1b.databricks.sandbox.Command\x12\x17\n\x0fnext_page_token\x18\x02 \x01(\t\"j\n\x07\x43ommand\x12\x12\n\ncommand_id\x18\x01 \x01(\t\x12\x0b\n\x03\x63md\x18\x02 \x01(\t\x12\x0c\n\x04\x61rgs\x18\x03 \x03(\t\x12\x0b\n\x03pid\x18\x04 \x01(\x03\x12\x10\n\x08\x66inished\x18\x05 \x01(\x08\x12\x11\n\texit_code\x18\x06 \x01(\x05\"i\n\x15\x41ttachCommandResponse\x12\x37\n\rprocess_event\x18\x01 \x01(\x0b\x32 .databricks.sandbox.ProcessEvent\x12\x17\n\x0fsequence_number\x18\x02 \x01(\x03\"j\n\x16\x45xecuteCommandResponse\x12\x37\n\rprocess_event\x18\x01 \x01(\x0b\x32 .databricks.sandbox.ProcessEvent\x12\x17\n\x0fsequence_number\x18\x02 \x01(\x03\"\xa4\x01\n\x0cProcessEvent\x12/\n\x05start\x18\x01 \x01(\x0b\x32\x1e.databricks.sandbox.StartEventH\x00\x12-\n\x04\x64\x61ta\x18\x02 \x01(\x0b\x32\x1d.databricks.sandbox.DataEventH\x00\x12+\n\x03\x65nd\x18\x03 \x01(\x0b\x32\x1c.databricks.sandbox.EndEventH\x00\x42\x07\n\x05\x65vent\"-\n\nStartEvent\x12\x0b\n\x03pid\x18\x01 \x01(\x03\x12\x12\n\ncommand_id\x18\x02 \x01(\t\"H\n\tDataEvent\x12\x10\n\x06stdout\x18\x01 \x01(\x0cH\x00\x12\x10\n\x06stderr\x18\x02 \x01(\x0cH\x00\x12\r\n\x03pty\x18\x03 \x01(\x0cH\x00\x42\x08\n\x06output\"\x1d\n\x08\x45ndEvent\x12\x11\n\texit_code\x18\x01 \x01(\x05\"\'\n\tPtyConfig\x12\x0c\n\x04\x63ols\x18\x01 \x01(\x05\x12\x0c\n\x04rows\x18\x02 \x01(\x05\x32\xb0\x03\n\x0eSandboxService\x12k\n\x0e\x45xecuteCommand\x12).databricks.sandbox.ExecuteCommandRequest\x1a*.databricks.sandbox.ExecuteCommandResponse\"\x00\x30\x01\x12h\n\rAttachCommand\x12(.databricks.sandbox.AttachCommandRequest\x1a).databricks.sandbox.AttachCommandResponse\"\x00\x30\x01\x12\x62\n\x0bStreamInput\x12&.databricks.sandbox.StreamInputRequest\x1a\'.databricks.sandbox.StreamInputResponse\"\x00(\x01\x12\x63\n\x0cListCommands\x12\'.databricks.sandbox.ListCommandsRequest\x1a(.databricks.sandbox.ListCommandsResponse\"\x00') + +_globals = globals() +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'databricks.sdk._sandbox_grpc.sandbox_service_pb2', _globals) +if not _descriptor._USE_C_DESCRIPTORS: + DESCRIPTOR._loaded_options = None + _globals['_EXECUTECOMMANDREQUEST_ENVSENTRY']._loaded_options = None + _globals['_EXECUTECOMMANDREQUEST_ENVSENTRY']._serialized_options = b'8\001' + _globals['_EXECUTECOMMANDREQUEST']._serialized_start=75 + _globals['_EXECUTECOMMANDREQUEST']._serialized_end=301 + _globals['_EXECUTECOMMANDREQUEST_ENVSENTRY']._serialized_start=258 + _globals['_EXECUTECOMMANDREQUEST_ENVSENTRY']._serialized_end=301 + _globals['_ATTACHCOMMANDREQUEST']._serialized_start=303 + _globals['_ATTACHCOMMANDREQUEST']._serialized_end=375 + _globals['_STREAMINPUTREQUEST']._serialized_start=378 + _globals['_STREAMINPUTREQUEST']._serialized_end=570 + _globals['_STREAMINPUTSTART']._serialized_start=572 + _globals['_STREAMINPUTSTART']._serialized_end=610 + _globals['_STREAMINPUTDATA']._serialized_start=612 + _globals['_STREAMINPUTDATA']._serialized_end=643 + _globals['_STREAMINPUTCLOSE']._serialized_start=645 + _globals['_STREAMINPUTCLOSE']._serialized_end=663 + _globals['_STREAMINPUTRESPONSE']._serialized_start=665 + _globals['_STREAMINPUTRESPONSE']._serialized_end=686 + _globals['_LISTCOMMANDSREQUEST']._serialized_start=688 + _globals['_LISTCOMMANDSREQUEST']._serialized_end=748 + _globals['_LISTCOMMANDSRESPONSE']._serialized_start=750 + _globals['_LISTCOMMANDSRESPONSE']._serialized_end=844 + _globals['_COMMAND']._serialized_start=846 + _globals['_COMMAND']._serialized_end=952 + _globals['_ATTACHCOMMANDRESPONSE']._serialized_start=954 + _globals['_ATTACHCOMMANDRESPONSE']._serialized_end=1059 + _globals['_EXECUTECOMMANDRESPONSE']._serialized_start=1061 + _globals['_EXECUTECOMMANDRESPONSE']._serialized_end=1167 + _globals['_PROCESSEVENT']._serialized_start=1170 + _globals['_PROCESSEVENT']._serialized_end=1334 + _globals['_STARTEVENT']._serialized_start=1336 + _globals['_STARTEVENT']._serialized_end=1381 + _globals['_DATAEVENT']._serialized_start=1383 + _globals['_DATAEVENT']._serialized_end=1455 + _globals['_ENDEVENT']._serialized_start=1457 + _globals['_ENDEVENT']._serialized_end=1486 + _globals['_PTYCONFIG']._serialized_start=1488 + _globals['_PTYCONFIG']._serialized_end=1527 + _globals['_SANDBOXSERVICE']._serialized_start=1530 + _globals['_SANDBOXSERVICE']._serialized_end=1962 +# @@protoc_insertion_point(module_scope) diff --git a/databricks/sdk/_sandbox_grpc/sandbox_service_pb2_grpc.py b/databricks/sdk/_sandbox_grpc/sandbox_service_pb2_grpc.py new file mode 100644 index 000000000..2889cc8fe --- /dev/null +++ b/databricks/sdk/_sandbox_grpc/sandbox_service_pb2_grpc.py @@ -0,0 +1,233 @@ +# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! +"""Client and server classes corresponding to protobuf-defined services.""" +import grpc +import warnings + +from databricks.sdk._sandbox_grpc import sandbox_service_pb2 as databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2 + +GRPC_GENERATED_VERSION = '1.81.1' +GRPC_VERSION = grpc.__version__ +_version_not_supported = False + +try: + from grpc._utilities import first_version_is_lower + _version_not_supported = first_version_is_lower(GRPC_VERSION, GRPC_GENERATED_VERSION) +except ImportError: + _version_not_supported = True + +if _version_not_supported: + raise RuntimeError( + f'The grpc package installed is at version {GRPC_VERSION},' + + ' but the generated code in databricks/sdk/_sandbox_grpc/sandbox_service_pb2_grpc.py depends on' + + f' grpcio>={GRPC_GENERATED_VERSION}.' + + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}' + + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.' + ) + + +class SandboxServiceStub: + """Service for executing commands inside running sandboxes. + """ + + def __init__(self, channel): + """Constructor. + + Args: + channel: A grpc.Channel. + """ + self.ExecuteCommand = channel.unary_stream( + '/databricks.sandbox.SandboxService/ExecuteCommand', + request_serializer=databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.ExecuteCommandRequest.SerializeToString, + response_deserializer=databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.ExecuteCommandResponse.FromString, + _registered_method=True) + self.AttachCommand = channel.unary_stream( + '/databricks.sandbox.SandboxService/AttachCommand', + request_serializer=databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.AttachCommandRequest.SerializeToString, + response_deserializer=databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.AttachCommandResponse.FromString, + _registered_method=True) + self.StreamInput = channel.stream_unary( + '/databricks.sandbox.SandboxService/StreamInput', + request_serializer=databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.StreamInputRequest.SerializeToString, + response_deserializer=databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.StreamInputResponse.FromString, + _registered_method=True) + self.ListCommands = channel.unary_unary( + '/databricks.sandbox.SandboxService/ListCommands', + request_serializer=databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.ListCommandsRequest.SerializeToString, + response_deserializer=databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.ListCommandsResponse.FromString, + _registered_method=True) + + +class SandboxServiceServicer: + """Service for executing commands inside running sandboxes. + """ + + def ExecuteCommand(self, request, context): + """Executes a command inside the sandbox. Streams process lifecycle events back to the caller. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def AttachCommand(self, request, context): + """Attaches to a previously started command execution, streaming its output. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def StreamInput(self, request_iterator, context): + """Streams stdin data to a running command. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def ListCommands(self, request, context): + """Lists all tracked command executions (running and completed). + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + +def add_SandboxServiceServicer_to_server(servicer, server): + rpc_method_handlers = { + 'ExecuteCommand': grpc.unary_stream_rpc_method_handler( + servicer.ExecuteCommand, + request_deserializer=databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.ExecuteCommandRequest.FromString, + response_serializer=databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.ExecuteCommandResponse.SerializeToString, + ), + 'AttachCommand': grpc.unary_stream_rpc_method_handler( + servicer.AttachCommand, + request_deserializer=databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.AttachCommandRequest.FromString, + response_serializer=databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.AttachCommandResponse.SerializeToString, + ), + 'StreamInput': grpc.stream_unary_rpc_method_handler( + servicer.StreamInput, + request_deserializer=databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.StreamInputRequest.FromString, + response_serializer=databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.StreamInputResponse.SerializeToString, + ), + 'ListCommands': grpc.unary_unary_rpc_method_handler( + servicer.ListCommands, + request_deserializer=databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.ListCommandsRequest.FromString, + response_serializer=databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.ListCommandsResponse.SerializeToString, + ), + } + generic_handler = grpc.method_handlers_generic_handler( + 'databricks.sandbox.SandboxService', rpc_method_handlers) + server.add_generic_rpc_handlers((generic_handler,)) + server.add_registered_method_handlers('databricks.sandbox.SandboxService', rpc_method_handlers) + + + # This class is part of an EXPERIMENTAL API. +class SandboxService: + """Service for executing commands inside running sandboxes. + """ + + @staticmethod + def ExecuteCommand(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_stream( + request, + target, + '/databricks.sandbox.SandboxService/ExecuteCommand', + databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.ExecuteCommandRequest.SerializeToString, + databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.ExecuteCommandResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def AttachCommand(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_stream( + request, + target, + '/databricks.sandbox.SandboxService/AttachCommand', + databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.AttachCommandRequest.SerializeToString, + databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.AttachCommandResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def StreamInput(request_iterator, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.stream_unary( + request_iterator, + target, + '/databricks.sandbox.SandboxService/StreamInput', + databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.StreamInputRequest.SerializeToString, + databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.StreamInputResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def ListCommands(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/databricks.sandbox.SandboxService/ListCommands', + databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.ListCommandsRequest.SerializeToString, + databricks_dot_sdk_dot___sandbox__grpc_dot_sandbox__service__pb2.ListCommandsResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) diff --git a/databricks/sdk/mixins/sandbox.py b/databricks/sdk/mixins/sandbox.py new file mode 100644 index 000000000..8b047a72a --- /dev/null +++ b/databricks/sdk/mixins/sandbox.py @@ -0,0 +1,187 @@ +"""Sandbox command-execution mixin. + +Extends the generated :class:`~databricks.sdk.service.sandbox.SandboxAPI` with +the streaming command RPCs that codegen does not produce a client for: +``execute_command_streaming``, ``attach_command`` (both gRPC server-streaming) +and ``stream_input`` (gRPC client-streaming). They run on the generic gRPC +transport in :mod:`databricks.sdk.mixins._grpc_transport`. + +The wire types come from the committed stubs in +:mod:`databricks.sdk._sandbox_grpc` (generated from a slim copy of the sandbox +daemon proto; see that package for provenance). grpcio and those stubs are +imported lazily inside the methods, so this module imports without the optional +``sandbox`` extra installed. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Callable, Dict, Iterable, Iterator, Optional, Tuple + +from databricks.sdk.mixins._grpc_transport import ( + _DEFAULT_MAX_ATTEMPTS, + _DEFAULT_RETRY_INTERVAL_SECONDS, + open_channel, + open_stream_with_retry, +) +from databricks.sdk.service.sandbox import SandboxAPI + +if TYPE_CHECKING: + from databricks.sdk._sandbox_grpc import sandbox_service_pb2 as _pb + from databricks.sdk._sandbox_grpc import sandbox_service_pb2_grpc as _pb_grpc + +# gRPC metadata key that routes a command RPC to a specific sandbox. +_SANDBOX_NAME_METADATA = "x-databricks-sandbox-name" + +_Metadata = Tuple[Tuple[str, str], ...] + + +def _never_retry(_error: Exception) -> bool: + """Retry predicate for :meth:`SandboxExt.execute_command_streaming`: never + retry. ExecuteCommand starts a process and is not idempotent (its + ``command_id`` is a stream-resume token, not a dedup key), so re-issuing a + failed open could run the command twice.""" + return False + + +class SandboxExt(SandboxAPI): + __doc__ = SandboxAPI.__doc__ + + def execute_command_streaming( + self, + name: str, + cmd: str, + *, + args: Optional[list[str]] = None, + envs: Optional[Dict[str, str]] = None, + is_retryable: Optional[Callable[[Exception], bool]] = None, + max_attempts: int = _DEFAULT_MAX_ATTEMPTS, + retry_interval_seconds: float = _DEFAULT_RETRY_INTERVAL_SECONDS, + ) -> Iterator["_pb.ExecuteCommandResponse"]: + """Run ``cmd`` in sandbox ``name`` and stream its process events. + + Streaming counterpart of :meth:`SandboxAPI.execute_command_sync`: the same + ``name`` (``sandboxes/{sandbox_id}``), ``cmd``, ``args``, and ``envs``, + but the process output is streamed instead of returned in one response. + Yields ``ExecuteCommandResponse`` messages, each carrying one + ``ProcessEvent`` (a start event, then stdout/stderr data chunks, then an + end event with the exit code) and a monotonic ``sequence_number``. + + Retries are off by default (``is_retryable`` defaults to + :func:`_never_retry`) because ExecuteCommand is not idempotent: re-issuing + a failed open could run the command twice. Pass an ``is_retryable`` + predicate - e.g. one matching ``UNAVAILABLE`` - to retry the stream open + when the command is safe to run more than once, optionally tuning + ``max_attempts`` and ``retry_interval_seconds``. Only the stream open is + retried; a failure once events are flowing is always raised. + """ + from databricks.sdk._sandbox_grpc import sandbox_service_pb2 as _pb # noqa: FlagLocalImports - lazy; stubs ship only with the optional "sandbox" extra + + request = _pb.ExecuteCommandRequest(cmd=cmd, args=list(args or []), envs=dict(envs or {})) + yield from self._server_stream( + name, + lambda stub, metadata: stub.ExecuteCommand(request, metadata=metadata), + is_retryable=_never_retry if is_retryable is None else is_retryable, + max_attempts=max_attempts, + retry_interval_seconds=retry_interval_seconds, + ) + + def attach_command( + self, + name: str, + command_id: str, + *, + last_sequence_number: Optional[int] = None, + is_retryable: Optional[Callable[[Exception], bool]] = None, + max_attempts: int = _DEFAULT_MAX_ATTEMPTS, + retry_interval_seconds: float = _DEFAULT_RETRY_INTERVAL_SECONDS, + ) -> Iterator["_pb.AttachCommandResponse"]: + """Attach to a running command in sandbox ``name`` and stream its output. + + ``command_id`` comes from the ``StartEvent`` of a prior + :meth:`execute_command_streaming`. The server replays buffered events with + ``sequence_number`` greater than ``last_sequence_number`` (0 replays from + the beginning; omitted tails live output only), then streams live events. + + Unlike :meth:`execute_command_streaming`, retries default to the + transport's ``UNAVAILABLE`` policy: attach is a resumable replay, so + re-opening from ``last_sequence_number`` is safe. Pass ``is_retryable`` to + override. + """ + from databricks.sdk._sandbox_grpc import sandbox_service_pb2 as _pb # noqa: FlagLocalImports - lazy; stubs ship only with the optional "sandbox" extra + + request = _pb.AttachCommandRequest(command_id=command_id) + if last_sequence_number is not None: + request.last_sequence_number = last_sequence_number + yield from self._server_stream( + name, + lambda stub, metadata: stub.AttachCommand(request, metadata=metadata), + is_retryable=is_retryable, + max_attempts=max_attempts, + retry_interval_seconds=retry_interval_seconds, + ) + + def stream_input( + self, + name: str, + command_id: str, + inputs: Iterable[bytes], + *, + close_stdin: bool = True, + ) -> "_pb.StreamInputResponse": + """Forward stdin ``inputs`` to a running command in sandbox ``name``. + + ``command_id`` comes from the ``StartEvent`` of a prior + :meth:`execute_command_streaming`. Sends a start message, one data message + per chunk of ``inputs``, then (unless ``close_stdin`` is false) a close + message that sends EOF to the process's stdin. Returns the server's + acknowledgement once the input stream is fully sent. + + Not retried: forwarding input is stateful and side-effecting, so a + transport failure is raised to the caller. + """ + from databricks.sdk._sandbox_grpc import sandbox_service_pb2 as _pb # noqa: FlagLocalImports - lazy; stubs ship only with the optional "sandbox" extra + from databricks.sdk._sandbox_grpc import sandbox_service_pb2_grpc as _pb_grpc # noqa: FlagLocalImports - lazy; stubs ship only with the optional "sandbox" extra + + def requests() -> Iterator["_pb.StreamInputRequest"]: + yield _pb.StreamInputRequest(start=_pb.StreamInputStart(command_id=command_id)) + for chunk in inputs: + yield _pb.StreamInputRequest(data=_pb.StreamInputData(data=chunk)) + if close_stdin: + yield _pb.StreamInputRequest(close=_pb.StreamInputClose()) + + channel = open_channel(self._api._cfg) + try: + stub = _pb_grpc.SandboxServiceStub(channel) + metadata: _Metadata = ((_SANDBOX_NAME_METADATA, name),) + return stub.StreamInput(requests(), metadata=metadata) + finally: + channel.close() + + def _server_stream( + self, + name: str, + make_stream: Callable[["_pb_grpc.SandboxServiceStub", _Metadata], Iterator], + *, + is_retryable: Optional[Callable[[Exception], bool]], + max_attempts: int, + retry_interval_seconds: float, + ) -> Iterator: + """Open a per-call authenticated channel, run ``make_stream`` on the stub + under :func:`open_stream_with_retry`, and yield its items, closing the + channel once iteration ends or is abandoned.""" + # Lazy import: grpcio and the generated stubs ship only with the optional + # "sandbox" extra, so importing this module must not require them. + from databricks.sdk._sandbox_grpc import sandbox_service_pb2_grpc as _pb_grpc # noqa: FlagLocalImports - lazy; stubs ship only with the optional "sandbox" extra + + channel = open_channel(self._api._cfg) + try: + stub = _pb_grpc.SandboxServiceStub(channel) + metadata: _Metadata = ((_SANDBOX_NAME_METADATA, name),) + yield from open_stream_with_retry( + lambda: make_stream(stub, metadata), + is_retryable=is_retryable, + max_attempts=max_attempts, + retry_interval_seconds=retry_interval_seconds, + ) + finally: + channel.close() diff --git a/databricks/sdk/service/aifunctions.py b/databricks/sdk/service/aifunctions.py index bc7b25568..c1338b840 100644 --- a/databricks/sdk/service/aifunctions.py +++ b/databricks/sdk/service/aifunctions.py @@ -145,6 +145,156 @@ def from_dict(cls, d: Dict[str, Any]) -> AiClassifyResponseMetadata: return cls(version=d.get("version", None)) +@dataclass +class AiDecideOptions: + """ai_decide""" + + version: Optional[str] = None + """The function API version to invoke. Defaults to "1.0". Supported versions: ["1.0"].""" + + def as_dict(self) -> dict: + """Serializes the AiDecideOptions into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.version is not None: + body["version"] = self.version + return body + + def as_shallow_dict(self) -> dict: + """Serializes the AiDecideOptions into a shallow dictionary of its immediate attributes.""" + body = {} + if self.version is not None: + body["version"] = self.version + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> AiDecideOptions: + """Deserializes the AiDecideOptions from a dictionary.""" + return cls(version=d.get("version", None)) + + +@dataclass +class AiDecideResponse: + metadata: Optional[AiDecideResponseMetadata] = None + """Metadata identifying the function API version used for this request.""" + + response: Optional[any] = None + """A JSON object containing an ``answers`` map with one answer for each question, keyed by the same + IDs as ``questions``. + + Each answer has a ``type`` matching its question and the following fields: + + - ``choice``: choice is the highest-probability option from ``criteria``. ``probabilities`` maps + every option name to its probability. ``confidence`` is a number from 0 to 1 indicating how + well the state supports the assessment. The ``probabilities`` values sum to 1. For example: + + .. code-block:: json + + { + "answers": { + "team": { + "type": "choice", + "choice": "billing", + "probabilities": { + "billing": 0.85, + "technical_support": 0.15 + }, + "confidence": 0.9 + } + } + } + + - ``noul``: ``probability`` is a number from 0 to 1 estimating the probability that the answer + is true. For example: + + .. code-block:: json + + { + "answers": { + "escalate": { + "type": "noul", + "probability": 0.8 + } + } + } + + - ``score``: score is the probability-weighted mean of the zero-based level indices. It can fall + between levels, from 0 to the number of levels minus 1. ``probabilities`` maps each level + index to its probability, and ``legend`` maps each index to its original description. Both + maps use string keys such as "0", "1", and "2". ``confidence`` is a number from 0 to 1 + indicating how well the state supports the assessment. The ``probabilities`` values sum to 1. + For example: + + .. code-block:: json + + { + "answers": { + "urgency": { + "type": "score", + "score": 1.6, + "probabilities": { + "0": 0.1, + "1": 0.2, + "2": 0.7 + }, + "legend": { + "0": "Routine: can wait a few days", + "1": "Time-sensitive: needs attention today", + "2": "Critical: needs immediate action" + }, + "confidence": 0.85 + } + } + }""" + + def as_dict(self) -> dict: + """Serializes the AiDecideResponse into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.metadata: + body["metadata"] = self.metadata.as_dict() + if self.response: + body["response"] = self.response + return body + + def as_shallow_dict(self) -> dict: + """Serializes the AiDecideResponse into a shallow dictionary of its immediate attributes.""" + body = {} + if self.metadata: + body["metadata"] = self.metadata + if self.response: + body["response"] = self.response + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> AiDecideResponse: + """Deserializes the AiDecideResponse from a dictionary.""" + return cls(metadata=_from_dict(d, "metadata", AiDecideResponseMetadata), response=d.get("response", None)) + + +@dataclass +class AiDecideResponseMetadata: + version: Optional[str] = None + """The function API version used to evaluate the request.""" + + def as_dict(self) -> dict: + """Serializes the AiDecideResponseMetadata into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.version is not None: + body["version"] = self.version + return body + + def as_shallow_dict(self) -> dict: + """Serializes the AiDecideResponseMetadata into a shallow dictionary of its immediate attributes.""" + body = {} + if self.version is not None: + body["version"] = self.version + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> AiDecideResponseMetadata: + """Deserializes the AiDecideResponseMetadata from a dictionary.""" + return cls(version=d.get("version", None)) + + @dataclass class AiExtractBbox: """A bounding box on a source page; used by bbox-input citations.""" @@ -668,6 +818,115 @@ def ai_classify( res = self._api.do("POST", "/api/2.0/ai-functions/ai-classify", body=body, headers=headers) return AiClassifyResponse.from_dict(res) + def ai_decide(self, state: any, questions: any, *, options: Optional[AiDecideOptions] = None) -> AiDecideResponse: + """Turn text and structured data into decisions your application can use. Define questions and criteria + to choose an option, estimate a probability, or assign a score given a provided state. + + :param state: any + A string, JSON object, or array containing the content, related context, and examples needed to + answer the provided questions. For example, provide a support message, a conversation, or records + describing the current state of an application. All questions receive this same state. + :param questions: any + A JSON object mapping question IDs to their definitions. Choose a nonempty string for each ID; its + answer is returned with the same ID in ``response.answers``. + + Each definition is an object with the required fields ``type`` and ``instructions``. The + ``criteria`` field is optional for the type ``noul`` but is required for the types ``choice`` and + ``score``. + + The ``instructions`` field describes the judgment to make and can be a string, object, or array. Use + an object or array to include supporting context alongside the instructions. + + The ``type`` can be one of: + + - ``choice``: Selects one option from a defined set. Requires ``criteria`` to be an object mapping 1 + to 255 nonempty option names to descriptions. The criteria description can be a string, object, + array, or null when the name needs no additional detail. For example: + + .. code-block:: json + + { + "team": { + "type": "choice", + "instructions": "Which team should handle this ticket?", + "criteria": { + "billing": "Payments, charges, and refunds", + "technical_support": null + } + } + } + + - ``noul``: Estimates the probability that the answer to a true-or-false question is true. + ``criteria`` can take the fields ``true`` or ``false``, or both, with descriptions that are + strings, objects, or arrays. Omit ``criteria`` to use the question alone. For example, both of the + following are valid: + + .. code-block:: json + + { + "escalate": { + "type": "noul", + "instructions": "Does this ticket need escalation?", + "criteria": { + "true": "Suspected fraud or an exception to standard policy", + "false": "A routine issue frontline support can resolve" + } + } + } + + or + + .. code-block:: json + + { + "escalate": { + "type": "noul", + "instructions": "Does this ticket need escalation?" + } + } + + - ``score``: Rates the state on an ordered scale. Requires ``criteria`` to be an array of 2 to 10 + level descriptions, ordered from low to high. Descriptions can be strings, objects, or arrays. + Array positions define levels starting at 0. For example: + + .. code-block:: json + + { + "urgency": { + "type": "score", + "instructions": "How urgent is this ticket?", + "criteria": [ + "Routine: can wait a few days", + "Time-sensitive: needs attention today", + "Critical: needs immediate action" + ] + } + } + :param options: :class:`AiDecideOptions` (optional) + Function options. Omitted fields fall back to their documented defaults. + + :returns: :class:`AiDecideResponse` + """ + + body = {} + if options is not None: + body["options"] = options.as_dict() + if questions is not None: + body["questions"] = questions + if state is not None: + body["state"] = state + headers = { + "Accept": "application/json", + "Content-Type": "application/json", + } + + cfg = self._api._cfg + if cfg.workspace_id: + headers["X-Databricks-Workspace-Id"] = cfg.workspace_id + + res = self._api.do("POST", "/api/2.0/ai-functions/ai-decide", body=body, headers=headers) + return AiDecideResponse.from_dict(res) + def ai_extract(self, content: any, schema: any, *, options: Optional[AiExtractOptions] = None) -> AiExtractResponse: """Extracts structured data from text and documents according to a provided schema. For REST API requests, the default rate limit is 120 requests per minute per workspace. Contact your Databricks diff --git a/databricks/sdk/service/compute.py b/databricks/sdk/service/compute.py index 5a07f08c9..c78cad91a 100644 --- a/databricks/sdk/service/compute.py +++ b/databricks/sdk/service/compute.py @@ -5455,7 +5455,12 @@ class InstancePoolAwsAttributes: """Attributes set during instance pool creation which are related to Amazon Web Services.""" availability: Optional[InstancePoolAwsAttributesAvailability] = None - """Availability type used for the spot nodes.""" + """Availability type used for the instances in the pool. Supports on-demand, spot, and + spot-with-fallback (the pool acquires spot instances first, and falls back to on-demand + instances when spot capacity is unavailable). + + You can change this value on an existing pool. New clusters use the updated availability, and + existing clusters keep the availability they launched with.""" instance_profile_arn: Optional[str] = None """All AWS instances belonging to the instance pool will have this instance profile. If omitted, @@ -5474,7 +5479,10 @@ class InstancePoolAwsAttributes: instances. Similarly, if this field is set to 200, the bid price is twice the price of on-demand ``r3.xlarge`` instances. If not specified, the default value is 100. When spot instances are requested for this cluster, only spot instances whose bid price percentage matches this field - will be considered. Note that, for safety, we enforce this field to be no more than 10000.""" + will be considered. Note that, for safety, we enforce this field to be no more than 10000. + + You can change this value on an existing pool. New clusters use the updated bid price, and + existing clusters keep the bid price they launched with.""" zone_id: Optional[str] = None """Identifier for the availability zone/datacenter in which the cluster resides. This string will @@ -5482,7 +5490,15 @@ class InstancePoolAwsAttributes: Databricks deployment. For example, "us-west-2a" is not a valid zone id if the Databricks deployment resides in the "us-east-1" region. This is an optional field at cluster creation, and if not specified, a default zone will be used. The list of available zones as well as the - default value can be found by using the ``List Zones`` method.""" + default value can be found by using the ``List Zones`` method. + + Set this field to "auto" to enable Auto-AZ, in which case Databricks selects the availability + zone for each cluster independently when the cluster launches, and retries another zone if the + cluster can't be fulfilled because of insufficient capacity or quota. All nodes in a cluster + land in the same zone, and different clusters backed by the pool can run in different zones. + + You can change the zone on an existing pool. New clusters use the updated zone, and existing + clusters keep the zone they launched with.""" def as_dict(self) -> dict: """Serializes the InstancePoolAwsAttributes into a dictionary suitable for use as a JSON request body.""" @@ -5533,7 +5549,12 @@ class InstancePoolAzureAttributes: """Attributes set during instance pool creation which are related to Azure.""" availability: Optional[InstancePoolAzureAttributesAvailability] = None - """Availability type used for the spot nodes.""" + """Availability type used for the instances in the pool. Supports on-demand, spot, and + spot-with-fallback (the pool acquires spot instances first, and falls back to on-demand + instances when spot capacity is unavailable). + + You can change this value on an existing pool. New clusters use the updated availability, and + existing clusters keep the availability they launched with.""" capacity_reservation_group: Optional[str] = None """The Azure capacity reservation group resource ID to use for launching VMs in this pool. When @@ -5558,7 +5579,10 @@ class InstancePoolAzureAttributes: """With variable pricing, you have option to set a max price, in US dollars (USD) For example, the value 2 would be a max price of $2.00 USD per hour. If you set the max price to be -1, the VM won't be evicted based on price. The price for the VM will be the current price for spot or the - price for a standard VM, which ever is less, as long as there is capacity and quota available.""" + price for a standard VM, which ever is less, as long as there is capacity and quota available. + + You can change this value on an existing pool. New clusters use the updated max price, and + existing clusters keep the max price they launched with.""" def as_dict(self) -> dict: """Serializes the InstancePoolAzureAttributes into a dictionary suitable for use as a JSON request body.""" @@ -5604,6 +5628,12 @@ class InstancePoolGcpAttributes: """Attributes set during instance pool creation which are related to GCP.""" gcp_availability: Optional[GcpAvailability] = None + """Availability type for the instances in the pool. One of: + + - ``ON_DEMAND_GCP``: the pool uses on-demand instances only. + - ``PREEMPTIBLE_GCP``: the pool uses preemptible instances only. + - ``PREEMPTIBLE_WITH_FALLBACK_GCP``: the pool acquires preemptible instances first, and falls + back to on-demand instances when preemptible capacity is unavailable.""" local_ssd_count: Optional[int] = None """If provided, each node in the instance pool will have this number of local SSDs attached. Each @@ -5622,10 +5652,17 @@ class InstancePoolGcpAttributes: - "HA" => High availability, spread nodes across availability zones for a Databricks deployment region + - "auto" => Auto-AZ. Databricks selects the availability zone for each cluster independently + when the cluster launches, and retries another zone if the cluster can't be fulfilled because + of insufficient capacity or quota. All nodes in a cluster land in the same zone, and different + clusters backed by the pool can run in different zones. - A GCP availability zone => Pick One of the available zones for (machine type + region) from https://cloud.google.com/compute/docs/regions-zones (e.g. "us-west1-a"). - If empty, Databricks picks an availability zone to schedule the cluster on.""" + If empty, Databricks picks an availability zone to schedule the cluster on. + + You can change the zone on an existing pool. New clusters use the updated zone, and existing + clusters keep the zone they launched with.""" def as_dict(self) -> dict: """Serializes the InstancePoolGcpAttributes into a dictionary suitable for use as a JSON request body.""" diff --git a/databricks/sdk/service/dashboards.py b/databricks/sdk/service/dashboards.py index 38cca7979..3e2c14f93 100644 --- a/databricks/sdk/service/dashboards.py +++ b/databricks/sdk/service/dashboards.py @@ -2901,7 +2901,7 @@ def delete_conversation_message(self, space_id: str, conversation_id: str, messa def download_message_attachment_visualization(self, name: str) -> DownloadMessageAttachmentVisualizationResponse: """Download a rendered image of a message visualization attachment. The response body is the raw PNG image, not a JSON payload. This is only available if the attachment is a visualization and the message - status is ``COMPLETED``. This endpoint is not supported for Private Link workspaces. + status is ``COMPLETED``. :param name: str The resource name of the attachment to render, in the format diff --git a/databricks/sdk/service/jobs.py b/databricks/sdk/service/jobs.py index e1fba90ad..d299b0a4a 100644 --- a/databricks/sdk/service/jobs.py +++ b/databricks/sdk/service/jobs.py @@ -513,6 +513,11 @@ class BaseRun: """The time at which this run ended in epoch milliseconds (milliseconds since 1/1/1970 UTC). This field is set to 0 if the job is still running.""" + environment_variables: Optional[List[JobEnvironmentVariables]] = None + """Snapshot of ``JobSettings.environment_variables`` as it was at run launch — the full list of + named environment-variable entries the job defined. To find which entry a given task ran with, + look at ``RunTaskSettings.environment_variables_key``.""" + execution_duration: Optional[int] = None """The time in milliseconds it took to execute the commands in the JAR or notebook until they completed, failed, timed out, were cancelled, or encountered an unexpected error. The duration @@ -640,6 +645,8 @@ def as_dict(self) -> dict: body["effective_usage_policy_id"] = self.effective_usage_policy_id if self.end_time is not None: body["end_time"] = self.end_time + if self.environment_variables: + body["environment_variables"] = [v.as_dict() for v in self.environment_variables] if self.execution_duration is not None: body["execution_duration"] = self.execution_duration if self.git_source: @@ -717,6 +724,8 @@ def as_shallow_dict(self) -> dict: body["effective_usage_policy_id"] = self.effective_usage_policy_id if self.end_time is not None: body["end_time"] = self.end_time + if self.environment_variables: + body["environment_variables"] = self.environment_variables if self.execution_duration is not None: body["execution_duration"] = self.execution_duration if self.git_source: @@ -785,6 +794,7 @@ def from_dict(cls, d: Dict[str, Any]) -> BaseRun: effective_performance_target=_enum(d, "effective_performance_target", PerformanceTarget), effective_usage_policy_id=d.get("effective_usage_policy_id", None), end_time=_int64(d, "end_time"), + environment_variables=_repeated_dict(d, "environment_variables", JobEnvironmentVariables), execution_duration=_int64(d, "execution_duration"), git_source=_from_dict(d, "git_source", GitSource), has_more=d.get("has_more", None), @@ -2103,7 +2113,10 @@ class DeploymentSpec: same command; role-split workloads (driver + worker, parameter server, separate eval node, etc.) use multiple entries.""" - command_path: str + compute: ComputeSpec + """Compute resources allocated to each node in this deployment.""" + + command_path: Optional[str] = None """Workspace path of the script to run on each node in this deployment. Upload the script to this path and supply the path here. When the task runs, the file at this path is run on each node; if it fails, the task fails with its exit code. @@ -2121,9 +2134,6 @@ class DeploymentSpec: # Distributed via torchrun: torchrun --nproc_per_node=8 train.py""" - compute: ComputeSpec - """Compute resources allocated to each node in this deployment.""" - name: Optional[str] = None """Optional human-readable name for this deployment (for example, ``driver``, ``worker``, ``param_server``). Used for log and UI display. Distinct names are recommended so deployments @@ -3380,6 +3390,100 @@ def from_dict(cls, d: Dict[str, Any]) -> JobEnvironment: return cls(environment_key=d.get("environment_key", None), spec=_from_dict(d, "spec", compute.Environment)) +@dataclass +class JobEnvironmentVariables: + """A named environment-variable entry, defined once at the job level and referenced by key from one + or more tasks. Entries live on ``JobSettings.environment_variables``, and tasks select one via + ``TaskSettings.environment_variables_key``.""" + + environment_variables_key: Optional[str] = None + """Identifier for this entry. Must be unique within ``JobSettings.environment_variables``. Tasks + reference it from ``TaskSettings.environment_variables_key``.""" + + spec: Optional[JobEnvironmentVariablesSpec] = None + """The environment variable specification.""" + + def as_dict(self) -> dict: + """Serializes the JobEnvironmentVariables into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.environment_variables_key is not None: + body["environment_variables_key"] = self.environment_variables_key + if self.spec: + body["spec"] = self.spec.as_dict() + return body + + def as_shallow_dict(self) -> dict: + """Serializes the JobEnvironmentVariables into a shallow dictionary of its immediate attributes.""" + body = {} + if self.environment_variables_key is not None: + body["environment_variables_key"] = self.environment_variables_key + if self.spec: + body["spec"] = self.spec + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> JobEnvironmentVariables: + """Deserializes the JobEnvironmentVariables from a dictionary.""" + return cls( + environment_variables_key=d.get("environment_variables_key", None), + spec=_from_dict(d, "spec", JobEnvironmentVariablesSpec), + ) + + +@dataclass +class JobEnvironmentVariablesSpec: + """The environment variables and files associated with a job environment variable entry. Runtime + environment variables override inline ``variables``, which override values from ``files``, on + duplicate keys.""" + + files: Optional[List[str]] = None + """Workspace (``/Workspace/...``) or UC Volumes (``/Volumes/...``) paths to ``.env`` files. Maximum + 5 files. Files are read, parsed, and merged at task execution time, not at job creation or + update API call time. + + File format: each line containing a variable must be exactly ``KEY=VALUE``. Empty and + whitespace-only lines, and lines beginning with ``#``, are ignored. Keys must match the same + regex as inlined variable names (``^[A-Za-z_][A-Za-z0-9_]*$``); the value continues to the end + of the line. No other syntax is supported — no inline comments, no quoted values, no escape + sequences, no variable interpolation. Any other line that does not match the ``KEY=VALUE`` shape + fails the run. + + Size limits: maximum 32,768 bytes (32 KiB) per file on disk; maximum 1,024 bytes (1 KiB) per + ``KEY=VALUE`` line combined. Files or lines exceeding these limits fail the run. + + On a duplicate key, the later file wins; ``variables`` override values from any file.""" + + variables: Optional[Dict[str, str]] = None + """Environment variables specified directly as key/value pairs. Maximum 20 entries. + + Each key must be 1 to 256 characters and match ``^[A-Za-z_][A-Za-z0-9_]*$``: it must start with + an ASCII letter or underscore and contain only ASCII letters, digits, and underscores. Each + value can be any Unicode string of up to 512 characters, including an empty string.""" + + def as_dict(self) -> dict: + """Serializes the JobEnvironmentVariablesSpec into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.files: + body["files"] = [v for v in self.files] + if self.variables: + body["variables"] = self.variables + return body + + def as_shallow_dict(self) -> dict: + """Serializes the JobEnvironmentVariablesSpec into a shallow dictionary of its immediate attributes.""" + body = {} + if self.files: + body["files"] = self.files + if self.variables: + body["variables"] = self.variables + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> JobEnvironmentVariablesSpec: + """Deserializes the JobEnvironmentVariablesSpec from a dictionary.""" + return cls(files=d.get("files", None), variables=d.get("variables", None)) + + @dataclass class JobNotificationSettings: no_alert_for_canceled_runs: Optional[bool] = None @@ -3693,6 +3797,12 @@ class JobSettings: """An optional set of email addresses that is notified when runs of this job begin or complete as well as when this job is deleted.""" + environment_variables: Optional[List[JobEnvironmentVariables]] = None + """Named environment-variable entries that tasks can reference by key from + ``TaskSettings.environment_variables_key``. Each entry's ``spec`` holds inline ``variables`` and + optional ``.env`` ``files``. Maximum 10 entries per job. A task can reference at most one entry + from this list.""" + environments: Optional[List[JobEnvironment]] = None """A list of task execution environment specifications that can be referenced by tasks that use serverless compute or a compute resource that uses Environments mode. @@ -3819,6 +3929,8 @@ def as_dict(self) -> dict: body["edit_mode"] = self.edit_mode.value if self.email_notifications: body["email_notifications"] = self.email_notifications.as_dict() + if self.environment_variables: + body["environment_variables"] = [v.as_dict() for v in self.environment_variables] if self.environments: body["environments"] = [v.as_dict() for v in self.environments] if self.format is not None: @@ -3878,6 +3990,8 @@ def as_shallow_dict(self) -> dict: body["edit_mode"] = self.edit_mode if self.email_notifications: body["email_notifications"] = self.email_notifications + if self.environment_variables: + body["environment_variables"] = self.environment_variables if self.environments: body["environments"] = self.environments if self.format is not None: @@ -3932,6 +4046,7 @@ def from_dict(cls, d: Dict[str, Any]) -> JobSettings: description=d.get("description", None), edit_mode=_enum(d, "edit_mode", JobEditMode), email_notifications=_from_dict(d, "email_notifications", JobEmailNotifications), + environment_variables=_repeated_dict(d, "environment_variables", JobEnvironmentVariables), environments=_repeated_dict(d, "environments", JobEnvironment), format=_enum(d, "format", Format), git_source=_from_dict(d, "git_source", GitSource), @@ -5738,6 +5853,11 @@ class Run: """The time at which this run ended in epoch milliseconds (milliseconds since 1/1/1970 UTC). This field is set to 0 if the job is still running.""" + environment_variables: Optional[List[JobEnvironmentVariables]] = None + """Snapshot of ``JobSettings.environment_variables`` as it was at run launch — the full list of + named environment-variable entries the job defined. To find which entry a given task ran with, + look at ``RunTaskSettings.environment_variables_key``.""" + execution_duration: Optional[int] = None """The time in milliseconds it took to execute the commands in the JAR or notebook until they completed, failed, timed out, were cancelled, or encountered an unexpected error. The duration @@ -5871,6 +5991,8 @@ def as_dict(self) -> dict: body["effective_usage_policy_id"] = self.effective_usage_policy_id if self.end_time is not None: body["end_time"] = self.end_time + if self.environment_variables: + body["environment_variables"] = [v.as_dict() for v in self.environment_variables] if self.execution_duration is not None: body["execution_duration"] = self.execution_duration if self.git_source: @@ -5952,6 +6074,8 @@ def as_shallow_dict(self) -> dict: body["effective_usage_policy_id"] = self.effective_usage_policy_id if self.end_time is not None: body["end_time"] = self.end_time + if self.environment_variables: + body["environment_variables"] = self.environment_variables if self.execution_duration is not None: body["execution_duration"] = self.execution_duration if self.git_source: @@ -6024,6 +6148,7 @@ def from_dict(cls, d: Dict[str, Any]) -> Run: effective_performance_target=_enum(d, "effective_performance_target", PerformanceTarget), effective_usage_policy_id=d.get("effective_usage_policy_id", None), end_time=_int64(d, "end_time"), + environment_variables=_repeated_dict(d, "environment_variables", JobEnvironmentVariables), execution_duration=_int64(d, "execution_duration"), git_source=_from_dict(d, "git_source", GitSource), has_more=d.get("has_more", None), @@ -6991,6 +7116,13 @@ class RunTask: Python wheel and dbt tasks when using serverless compute or a compute resource that uses Environments mode.""" + environment_variables_key: Optional[str] = None + """Reference to a ``JobEnvironmentVariables`` entry defined in + ``RunSettings.environment_variables`` for one-time runs or preserved in + ``Run.environment_variables`` for run snapshots. The selected entry's variables are applied to + this task at execution time. This field supports serverless tasks using environment version 5 or + later.""" + execution_duration: Optional[int] = None """The time in milliseconds it took to execute the commands in the JAR or notebook until they completed, failed, timed out, were cancelled, or encountered an unexpected error. The duration @@ -7169,6 +7301,8 @@ def as_dict(self) -> dict: body["end_time"] = self.end_time if self.environment_key is not None: body["environment_key"] = self.environment_key + if self.environment_variables_key is not None: + body["environment_variables_key"] = self.environment_variables_key if self.execution_duration is not None: body["execution_duration"] = self.execution_duration if self.existing_cluster_id is not None: @@ -7286,6 +7420,8 @@ def as_shallow_dict(self) -> dict: body["end_time"] = self.end_time if self.environment_key is not None: body["environment_key"] = self.environment_key + if self.environment_variables_key is not None: + body["environment_variables_key"] = self.environment_variables_key if self.execution_duration is not None: body["execution_duration"] = self.execution_duration if self.existing_cluster_id is not None: @@ -7383,6 +7519,7 @@ def from_dict(cls, d: Dict[str, Any]) -> RunTask: email_notifications=_from_dict(d, "email_notifications", JobEmailNotifications), end_time=_int64(d, "end_time"), environment_key=d.get("environment_key", None), + environment_variables_key=d.get("environment_variables_key", None), execution_duration=_int64(d, "execution_duration"), existing_cluster_id=d.get("existing_cluster_id", None), for_each_task=_from_dict(d, "for_each_task", RunForEachTask), @@ -8493,6 +8630,13 @@ class SubmitTask: Python wheel and dbt tasks when using serverless compute or a compute resource that uses Environments mode.""" + environment_variables_key: Optional[str] = None + """Reference to a ``JobEnvironmentVariables`` entry defined in + ``RunSettings.environment_variables`` for one-time runs or preserved in + ``Run.environment_variables`` for run snapshots. The selected entry's variables are applied to + this task at execution time. This field supports serverless tasks using environment version 5 or + later.""" + existing_cluster_id: Optional[str] = None """If existing_cluster_id, the ID of an existing cluster that is used for all runs. When running jobs or tasks on an existing cluster, you may need to manually restart the cluster if it stops @@ -8611,6 +8755,8 @@ def as_dict(self) -> dict: body["email_notifications"] = self.email_notifications.as_dict() if self.environment_key is not None: body["environment_key"] = self.environment_key + if self.environment_variables_key is not None: + body["environment_variables_key"] = self.environment_variables_key if self.existing_cluster_id is not None: body["existing_cluster_id"] = self.existing_cluster_id if self.for_each_task: @@ -8694,6 +8840,8 @@ def as_shallow_dict(self) -> dict: body["email_notifications"] = self.email_notifications if self.environment_key is not None: body["environment_key"] = self.environment_key + if self.environment_variables_key is not None: + body["environment_variables_key"] = self.environment_variables_key if self.existing_cluster_id is not None: body["existing_cluster_id"] = self.existing_cluster_id if self.for_each_task: @@ -8763,6 +8911,7 @@ def from_dict(cls, d: Dict[str, Any]) -> SubmitTask: disabled=d.get("disabled", None), email_notifications=_from_dict(d, "email_notifications", JobEmailNotifications), environment_key=d.get("environment_key", None), + environment_variables_key=d.get("environment_variables_key", None), existing_cluster_id=d.get("existing_cluster_id", None), for_each_task=_from_dict(d, "for_each_task", ForEachTask), gen_ai_compute_task=_from_dict(d, "gen_ai_compute_task", GenAiComputeTask), @@ -9053,6 +9202,11 @@ class Task: Python wheel and dbt tasks when using serverless compute or a compute resource that uses Environments mode.""" + environment_variables_key: Optional[str] = None + """Reference to a ``JobEnvironmentVariables`` entry defined in + ``JobSettings.environment_variables``. The selected entry's variables are applied to this task + at execution time. This field supports serverless tasks using environment version 5 or later.""" + existing_cluster_id: Optional[str] = None """If existing_cluster_id, the ID of an existing cluster that is used for all runs. When running jobs or tasks on an existing cluster, you may need to manually restart the cluster if it stops @@ -9180,6 +9334,8 @@ def as_dict(self) -> dict: body["email_notifications"] = self.email_notifications.as_dict() if self.environment_key is not None: body["environment_key"] = self.environment_key + if self.environment_variables_key is not None: + body["environment_variables_key"] = self.environment_variables_key if self.existing_cluster_id is not None: body["existing_cluster_id"] = self.existing_cluster_id if self.for_each_task: @@ -9265,6 +9421,8 @@ def as_shallow_dict(self) -> dict: body["email_notifications"] = self.email_notifications if self.environment_key is not None: body["environment_key"] = self.environment_key + if self.environment_variables_key is not None: + body["environment_variables_key"] = self.environment_variables_key if self.existing_cluster_id is not None: body["existing_cluster_id"] = self.existing_cluster_id if self.for_each_task: @@ -9336,6 +9494,7 @@ def from_dict(cls, d: Dict[str, Any]) -> Task: disabled=d.get("disabled", None), email_notifications=_from_dict(d, "email_notifications", TaskEmailNotifications), environment_key=d.get("environment_key", None), + environment_variables_key=d.get("environment_variables_key", None), existing_cluster_id=d.get("existing_cluster_id", None), for_each_task=_from_dict(d, "for_each_task", ForEachTask), gen_ai_compute_task=_from_dict(d, "gen_ai_compute_task", GenAiComputeTask), @@ -10406,6 +10565,7 @@ def create( description: Optional[str] = None, edit_mode: Optional[JobEditMode] = None, email_notifications: Optional[JobEmailNotifications] = None, + environment_variables: Optional[List[JobEnvironmentVariables]] = None, environments: Optional[List[JobEnvironment]] = None, format: Optional[Format] = None, git_source: Optional[GitSource] = None, @@ -10454,6 +10614,11 @@ def create( :param email_notifications: :class:`JobEmailNotifications` (optional) An optional set of email addresses that is notified when runs of this job begin or complete as well as when this job is deleted. + :param environment_variables: List[:class:`JobEnvironmentVariables`] (optional) + Named environment-variable entries that tasks can reference by key from + ``TaskSettings.environment_variables_key``. Each entry's ``spec`` holds inline ``variables`` and + optional ``.env`` ``files``. Maximum 10 entries per job. A task can reference at most one entry from + this list. :param environments: List[:class:`JobEnvironment`] (optional) A list of task execution environment specifications that can be referenced by tasks that use serverless compute or a compute resource that uses Environments mode. @@ -10562,6 +10727,8 @@ def create( body["edit_mode"] = edit_mode.value if email_notifications is not None: body["email_notifications"] = email_notifications.as_dict() + if environment_variables is not None: + body["environment_variables"] = [v.as_dict() for v in environment_variables] if environments is not None: body["environments"] = [v.as_dict() for v in environments] if format is not None: @@ -11480,6 +11647,7 @@ def submit( access_control_list: Optional[List[JobAccessControlRequest]] = None, budget_policy_id: Optional[str] = None, email_notifications: Optional[JobEmailNotifications] = None, + environment_variables: Optional[List[JobEnvironmentVariables]] = None, environments: Optional[List[JobEnvironment]] = None, git_source: Optional[GitSource] = None, health: Optional[JobsHealthRules] = None, @@ -11511,6 +11679,11 @@ def submit( will be not be attributed to any budget policy. :param email_notifications: :class:`JobEmailNotifications` (optional) An optional set of email addresses notified when the run begins or completes. + :param environment_variables: List[:class:`JobEnvironmentVariables`] (optional) + Named environment-variable entries that tasks of this one-time run can reference by key from + ``RunTaskSettings.environment_variables_key``. Each entry's ``spec`` holds inline ``variables`` and + optional ``.env`` ``files``. Handled identically to ``JobSettings.environment_variables``. Maximum + 10 entries. Entries are independent of one another — there is no cross-entry merging. :param environments: List[:class:`JobEnvironment`] (optional) A list of task execution environment specifications that can be referenced by tasks of this run. :param git_source: :class:`GitSource` (optional) @@ -11573,6 +11746,8 @@ def submit( body["budget_policy_id"] = budget_policy_id if email_notifications is not None: body["email_notifications"] = email_notifications.as_dict() + if environment_variables is not None: + body["environment_variables"] = [v.as_dict() for v in environment_variables] if environments is not None: body["environments"] = [v.as_dict() for v in environments] if git_source is not None: @@ -11621,6 +11796,7 @@ def submit_and_wait( access_control_list: Optional[List[JobAccessControlRequest]] = None, budget_policy_id: Optional[str] = None, email_notifications: Optional[JobEmailNotifications] = None, + environment_variables: Optional[List[JobEnvironmentVariables]] = None, environments: Optional[List[JobEnvironment]] = None, git_source: Optional[GitSource] = None, health: Optional[JobsHealthRules] = None, @@ -11640,6 +11816,7 @@ def submit_and_wait( access_control_list=access_control_list, budget_policy_id=budget_policy_id, email_notifications=email_notifications, + environment_variables=environment_variables, environments=environments, git_source=git_source, health=health, diff --git a/databricks/sdk/service/ml.py b/databricks/sdk/service/ml.py index a55c4982e..0ea1b331e 100644 --- a/databricks/sdk/service/ml.py +++ b/databricks/sdk/service/ml.py @@ -7713,8 +7713,9 @@ class SchemaLocatorFormat(Enum): class SchemaRegistryConfig: """Configuration for resolving a Stream's schema from an external schema registry (e.g. Confluent).""" - api_secret_ref: SecretScopeReference - """Reference to the schema registry API secret in a Databricks secret scope.""" + api_secret_ref: Optional[SecretScopeReference] = None + """Reference to the schema registry API secret in a Databricks secret scope. Set this only if + required for authentication for the schema registry.""" key_schema_locator: Optional[SchemaLocator] = None """Schema locator for the message key. Only used for Kafka streams. At least one of diff --git a/databricks/sdk/service/pipelines.py b/databricks/sdk/service/pipelines.py index 4c74c0a68..2bad4a7bf 100644 --- a/databricks/sdk/service/pipelines.py +++ b/databricks/sdk/service/pipelines.py @@ -3708,7 +3708,9 @@ class PipelineSpec: """Deployment type of this pipeline.""" development: Optional[bool] = None - """Whether the pipeline is in Development mode. Defaults to false.""" + """Whether the pipeline is in Development mode. Defaults to false. + + Deprecated: set development mode for each update instead.""" edition: Optional[str] = None """Pipeline product edition.""" @@ -6075,6 +6077,8 @@ def clone( Deployment type of this pipeline. :param development: bool (optional) Whether the pipeline is in Development mode. Defaults to false. + + Deprecated: set development mode for each update instead. :param edition: str (optional) Pipeline product edition. :param environment: :class:`PipelinesEnvironment` (optional) @@ -6272,6 +6276,8 @@ def create( Deployment type of this pipeline. :param development: bool (optional) Whether the pipeline is in Development mode. Defaults to false. + + Deprecated: set development mode for each update instead. :param dry_run: bool (optional) :param edition: str (optional) Pipeline product edition. @@ -6887,6 +6893,8 @@ def update( Deployment type of this pipeline. :param development: bool (optional) Whether the pipeline is in Development mode. Defaults to false. + + Deprecated: set development mode for each update instead. :param edition: str (optional) Pipeline product edition. :param environment: :class:`PipelinesEnvironment` (optional) diff --git a/databricks/sdk/version.py b/databricks/sdk/version.py index 5a6bb17fb..3b5229b1c 100644 --- a/databricks/sdk/version.py +++ b/databricks/sdk/version.py @@ -1 +1 @@ -__version__ = "0.143.0" +__version__ = "0.144.0" diff --git a/docs/dbdataclasses/aifunctions.rst b/docs/dbdataclasses/aifunctions.rst index 3f4f09d62..48bb9b4b0 100644 --- a/docs/dbdataclasses/aifunctions.rst +++ b/docs/dbdataclasses/aifunctions.rst @@ -16,6 +16,18 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: AiDecideOptions + :members: + :undoc-members: + +.. autoclass:: AiDecideResponse + :members: + :undoc-members: + +.. autoclass:: AiDecideResponseMetadata + :members: + :undoc-members: + .. autoclass:: AiExtractBbox :members: :undoc-members: diff --git a/docs/dbdataclasses/jobs.rst b/docs/dbdataclasses/jobs.rst index 6402a8ad6..6f6bccc03 100644 --- a/docs/dbdataclasses/jobs.rst +++ b/docs/dbdataclasses/jobs.rst @@ -491,6 +491,14 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: JobEnvironmentVariables + :members: + :undoc-members: + +.. autoclass:: JobEnvironmentVariablesSpec + :members: + :undoc-members: + .. autoclass:: JobNotificationSettings :members: :undoc-members: diff --git a/docs/workspace/aifunctions/ai_functions.rst b/docs/workspace/aifunctions/ai_functions.rst index 423f82ab4..f6fae9758 100644 --- a/docs/workspace/aifunctions/ai_functions.rst +++ b/docs/workspace/aifunctions/ai_functions.rst @@ -25,6 +25,97 @@ :returns: :class:`AiClassifyResponse` + .. py:method:: ai_decide(state: any, questions: any [, options: Optional[AiDecideOptions]]) -> AiDecideResponse + + Turn text and structured data into decisions your application can use. Define questions and criteria + to choose an option, estimate a probability, or assign a score given a provided state. + + :param state: any + A string, JSON object, or array containing the content, related context, and examples needed to + answer the provided questions. For example, provide a support message, a conversation, or records + describing the current state of an application. All questions receive this same state. + :param questions: any + A JSON object mapping question IDs to their definitions. Choose a nonempty string for each ID; its + answer is returned with the same ID in ``response.answers``. + + Each definition is an object with the required fields ``type`` and ``instructions``. The + ``criteria`` field is optional for the type ``noul`` but is required for the types ``choice`` and + ``score``. + + The ``instructions`` field describes the judgment to make and can be a string, object, or array. Use + an object or array to include supporting context alongside the instructions. + + The ``type`` can be one of: + + - ``choice``: Selects one option from a defined set. Requires ``criteria`` to be an object mapping 1 + to 255 nonempty option names to descriptions. The criteria description can be a string, object, + array, or null when the name needs no additional detail. For example: + + .. code-block:: json + + { + "team": { + "type": "choice", + "instructions": "Which team should handle this ticket?", + "criteria": { + "billing": "Payments, charges, and refunds", + "technical_support": null + } + } + } + + - ``noul``: Estimates the probability that the answer to a true-or-false question is true. + ``criteria`` can take the fields ``true`` or ``false``, or both, with descriptions that are + strings, objects, or arrays. Omit ``criteria`` to use the question alone. For example, both of the + following are valid: + + .. code-block:: json + + { + "escalate": { + "type": "noul", + "instructions": "Does this ticket need escalation?", + "criteria": { + "true": "Suspected fraud or an exception to standard policy", + "false": "A routine issue frontline support can resolve" + } + } + } + + or + + .. code-block:: json + + { + "escalate": { + "type": "noul", + "instructions": "Does this ticket need escalation?" + } + } + + - ``score``: Rates the state on an ordered scale. Requires ``criteria`` to be an array of 2 to 10 + level descriptions, ordered from low to high. Descriptions can be strings, objects, or arrays. + Array positions define levels starting at 0. For example: + + .. code-block:: json + + { + "urgency": { + "type": "score", + "instructions": "How urgent is this ticket?", + "criteria": [ + "Routine: can wait a few days", + "Time-sensitive: needs attention today", + "Critical: needs immediate action" + ] + } + } + :param options: :class:`AiDecideOptions` (optional) + Function options. Omitted fields fall back to their documented defaults. + + :returns: :class:`AiDecideResponse` + + .. py:method:: ai_extract(content: any, schema: any [, options: Optional[AiExtractOptions]]) -> AiExtractResponse Extracts structured data from text and documents according to a provided schema. For REST API diff --git a/docs/workspace/dashboards/genie.rst b/docs/workspace/dashboards/genie.rst index e2ac51e4a..3f47480ea 100644 --- a/docs/workspace/dashboards/genie.rst +++ b/docs/workspace/dashboards/genie.rst @@ -98,7 +98,7 @@ Download a rendered image of a message visualization attachment. The response body is the raw PNG image, not a JSON payload. This is only available if the attachment is a visualization and the message - status is ``COMPLETED``. This endpoint is not supported for Private Link workspaces. + status is ``COMPLETED``. :param name: str The resource name of the attachment to render, in the format diff --git a/docs/workspace/jobs/jobs.rst b/docs/workspace/jobs/jobs.rst index c4a51fbbe..10a217e02 100644 --- a/docs/workspace/jobs/jobs.rst +++ b/docs/workspace/jobs/jobs.rst @@ -125,7 +125,7 @@ .. py:method:: cancel_run_and_wait(run_id: int, timeout: datetime.timedelta = 0:20:00) -> Run - .. py:method:: create( [, access_control_list: Optional[List[JobAccessControlRequest]], budget_policy_id: Optional[str], continuous: Optional[Continuous], deployment: Optional[JobDeployment], description: Optional[str], edit_mode: Optional[JobEditMode], email_notifications: Optional[JobEmailNotifications], environments: Optional[List[JobEnvironment]], format: Optional[Format], git_source: Optional[GitSource], health: Optional[JobsHealthRules], job_clusters: Optional[List[JobCluster]], max_concurrent_runs: Optional[int], name: Optional[str], notification_settings: Optional[JobNotificationSettings], parameters: Optional[List[JobParameterDefinition]], parent_path: Optional[str], performance_target: Optional[PerformanceTarget], queue: Optional[QueueSettings], run_as: Optional[JobRunAs], schedule: Optional[CronSchedule], tags: Optional[Dict[str, str]], tasks: Optional[List[Task]], timeout_seconds: Optional[int], trigger: Optional[TriggerSettings], triggers: Optional[List[TriggerConfiguration]], usage_policy_id: Optional[str], webhook_notifications: Optional[WebhookNotifications]]) -> CreateResponse + .. py:method:: create( [, access_control_list: Optional[List[JobAccessControlRequest]], budget_policy_id: Optional[str], continuous: Optional[Continuous], deployment: Optional[JobDeployment], description: Optional[str], edit_mode: Optional[JobEditMode], email_notifications: Optional[JobEmailNotifications], environment_variables: Optional[List[JobEnvironmentVariables]], environments: Optional[List[JobEnvironment]], format: Optional[Format], git_source: Optional[GitSource], health: Optional[JobsHealthRules], job_clusters: Optional[List[JobCluster]], max_concurrent_runs: Optional[int], name: Optional[str], notification_settings: Optional[JobNotificationSettings], parameters: Optional[List[JobParameterDefinition]], parent_path: Optional[str], performance_target: Optional[PerformanceTarget], queue: Optional[QueueSettings], run_as: Optional[JobRunAs], schedule: Optional[CronSchedule], tags: Optional[Dict[str, str]], tasks: Optional[List[Task]], timeout_seconds: Optional[int], trigger: Optional[TriggerSettings], triggers: Optional[List[TriggerConfiguration]], usage_policy_id: Optional[str], webhook_notifications: Optional[WebhookNotifications]]) -> CreateResponse Usage: @@ -188,6 +188,11 @@ :param email_notifications: :class:`JobEmailNotifications` (optional) An optional set of email addresses that is notified when runs of this job begin or complete as well as when this job is deleted. + :param environment_variables: List[:class:`JobEnvironmentVariables`] (optional) + Named environment-variable entries that tasks can reference by key from + ``TaskSettings.environment_variables_key``. Each entry's ``spec`` holds inline ``variables`` and + optional ``.env`` ``files``. Maximum 10 entries per job. A task can reference at most one entry from + this list. :param environments: List[:class:`JobEnvironment`] (optional) A list of task execution environment specifications that can be referenced by tasks that use serverless compute or a compute resource that uses Environments mode. @@ -1030,7 +1035,7 @@ :returns: :class:`JobPermissions` - .. py:method:: submit( [, access_control_list: Optional[List[JobAccessControlRequest]], budget_policy_id: Optional[str], email_notifications: Optional[JobEmailNotifications], environments: Optional[List[JobEnvironment]], git_source: Optional[GitSource], health: Optional[JobsHealthRules], idempotency_token: Optional[str], notification_settings: Optional[JobNotificationSettings], performance_target: Optional[PerformanceTarget], queue: Optional[QueueSettings], run_as: Optional[JobRunAs], run_name: Optional[str], tasks: Optional[List[SubmitTask]], timeout_seconds: Optional[int], usage_policy_id: Optional[str], webhook_notifications: Optional[WebhookNotifications]]) -> Wait[Run] + .. py:method:: submit( [, access_control_list: Optional[List[JobAccessControlRequest]], budget_policy_id: Optional[str], email_notifications: Optional[JobEmailNotifications], environment_variables: Optional[List[JobEnvironmentVariables]], environments: Optional[List[JobEnvironment]], git_source: Optional[GitSource], health: Optional[JobsHealthRules], idempotency_token: Optional[str], notification_settings: Optional[JobNotificationSettings], performance_target: Optional[PerformanceTarget], queue: Optional[QueueSettings], run_as: Optional[JobRunAs], run_name: Optional[str], tasks: Optional[List[SubmitTask]], timeout_seconds: Optional[int], usage_policy_id: Optional[str], webhook_notifications: Optional[WebhookNotifications]]) -> Wait[Run] Usage: @@ -1082,6 +1087,11 @@ will be not be attributed to any budget policy. :param email_notifications: :class:`JobEmailNotifications` (optional) An optional set of email addresses notified when the run begins or completes. + :param environment_variables: List[:class:`JobEnvironmentVariables`] (optional) + Named environment-variable entries that tasks of this one-time run can reference by key from + ``RunTaskSettings.environment_variables_key``. Each entry's ``spec`` holds inline ``variables`` and + optional ``.env`` ``files``. Handled identically to ``JobSettings.environment_variables``. Maximum + 10 entries. Entries are independent of one another — there is no cross-entry merging. :param environments: List[:class:`JobEnvironment`] (optional) A list of task execution environment specifications that can be referenced by tasks of this run. :param git_source: :class:`GitSource` (optional) @@ -1137,7 +1147,7 @@ See :method:wait_get_run_job_terminated_or_skipped for more details. - .. py:method:: submit_and_wait( [, access_control_list: Optional[List[JobAccessControlRequest]], budget_policy_id: Optional[str], email_notifications: Optional[JobEmailNotifications], environments: Optional[List[JobEnvironment]], git_source: Optional[GitSource], health: Optional[JobsHealthRules], idempotency_token: Optional[str], notification_settings: Optional[JobNotificationSettings], performance_target: Optional[PerformanceTarget], queue: Optional[QueueSettings], run_as: Optional[JobRunAs], run_name: Optional[str], tasks: Optional[List[SubmitTask]], timeout_seconds: Optional[int], usage_policy_id: Optional[str], webhook_notifications: Optional[WebhookNotifications], timeout: datetime.timedelta = 0:20:00]) -> Run + .. py:method:: submit_and_wait( [, access_control_list: Optional[List[JobAccessControlRequest]], budget_policy_id: Optional[str], email_notifications: Optional[JobEmailNotifications], environment_variables: Optional[List[JobEnvironmentVariables]], environments: Optional[List[JobEnvironment]], git_source: Optional[GitSource], health: Optional[JobsHealthRules], idempotency_token: Optional[str], notification_settings: Optional[JobNotificationSettings], performance_target: Optional[PerformanceTarget], queue: Optional[QueueSettings], run_as: Optional[JobRunAs], run_name: Optional[str], tasks: Optional[List[SubmitTask]], timeout_seconds: Optional[int], usage_policy_id: Optional[str], webhook_notifications: Optional[WebhookNotifications], timeout: datetime.timedelta = 0:20:00]) -> Run .. py:method:: update(job_id: int [, fields_to_remove: Optional[List[str]], new_settings: Optional[JobSettings]]) diff --git a/docs/workspace/pipelines/pipelines.rst b/docs/workspace/pipelines/pipelines.rst index db6005427..6ee7e28de 100644 --- a/docs/workspace/pipelines/pipelines.rst +++ b/docs/workspace/pipelines/pipelines.rst @@ -61,6 +61,8 @@ Deployment type of this pipeline. :param development: bool (optional) Whether the pipeline is in Development mode. Defaults to false. + + Deprecated: set development mode for each update instead. :param edition: str (optional) Pipeline product edition. :param environment: :class:`PipelinesEnvironment` (optional) @@ -180,6 +182,8 @@ Deployment type of this pipeline. :param development: bool (optional) Whether the pipeline is in Development mode. Defaults to false. + + Deprecated: set development mode for each update instead. :param dry_run: bool (optional) :param edition: str (optional) Pipeline product edition. @@ -599,6 +603,8 @@ Deployment type of this pipeline. :param development: bool (optional) Whether the pipeline is in Development mode. Defaults to false. + + Deprecated: set development mode for each update instead. :param edition: str (optional) Pipeline product edition. :param environment: :class:`PipelinesEnvironment` (optional) diff --git a/docs/workspace/sandbox/sandbox.rst b/docs/workspace/sandbox/sandbox.rst index f686eedd2..00216a60c 100644 --- a/docs/workspace/sandbox/sandbox.rst +++ b/docs/workspace/sandbox/sandbox.rst @@ -2,11 +2,26 @@ ========================= .. currentmodule:: databricks.sdk.service.sandbox -.. py:class:: SandboxAPI +.. py:class:: SandboxExt Create, manage, and control the lifecycle of sandboxes -- isolated, pre-configured, low-latency Serverless compute environments for running code. + .. py:method:: attach_command(name: str, command_id: str [, last_sequence_number: Optional[int], is_retryable: Optional[Callable[[Exception], bool]], max_attempts: int = 30, retry_interval_seconds: float = 3.0]) -> Iterator['_pb.AttachCommandResponse'] + + Attach to a running command in sandbox ``name`` and stream its output. + + ``command_id`` comes from the ``StartEvent`` of a prior + :meth:`execute_command_streaming`. The server replays buffered events with + ``sequence_number`` greater than ``last_sequence_number`` (0 replays from + the beginning; omitted tails live output only), then streams live events. + + Unlike :meth:`execute_command_streaming`, retries default to the + transport's ``UNAVAILABLE`` policy: attach is a resumable replay, so + re-opening from ``last_sequence_number`` is safe. Pass ``is_retryable`` to + override. + + .. py:method:: create_sandbox(sandbox: Sandbox, sandbox_id: str) -> Sandbox Creates a new Sandbox. @@ -28,6 +43,26 @@ + .. py:method:: execute_command_streaming(name: str, cmd: str [, args: Optional[list[str]], envs: Optional[Dict[str, str]], is_retryable: Optional[Callable[[Exception], bool]], max_attempts: int = 30, retry_interval_seconds: float = 3.0]) -> Iterator['_pb.ExecuteCommandResponse'] + + Run ``cmd`` in sandbox ``name`` and stream its process events. + + Streaming counterpart of :meth:`SandboxAPI.execute_command_sync`: the same + ``name`` (``sandboxes/{sandbox_id}``), ``cmd``, ``args``, and ``envs``, + but the process output is streamed instead of returned in one response. + Yields ``ExecuteCommandResponse`` messages, each carrying one + ``ProcessEvent`` (a start event, then stdout/stderr data chunks, then an + end event with the exit code) and a monotonic ``sequence_number``. + + Retries are off by default (``is_retryable`` defaults to + :func:`_never_retry`) because ExecuteCommand is not idempotent: re-issuing + a failed open could run the command twice. Pass an ``is_retryable`` + predicate - e.g. one matching ``UNAVAILABLE`` - to retry the stream open + when the command is safe to run more than once, optionally tuning + ``max_attempts`` and ``retry_interval_seconds``. Only the stream open is + retried; a failure once events are flowing is always raised. + + .. py:method:: execute_command_sync(name: str, cmd: str [, args: Optional[List[str]], envs: Optional[Dict[str, str]], execution_timeout: Optional[Duration]]) -> ExecuteCommandSyncResponse Runs a command in the sandbox and blocks until it exits, returning the captured stdout, stderr and @@ -91,6 +126,20 @@ :returns: :class:`Sandbox` + .. py:method:: stream_input(name: str, command_id: str, inputs: Iterable[bytes] [, close_stdin: bool = True]) -> '_pb.StreamInputResponse' + + Forward stdin ``inputs`` to a running command in sandbox ``name``. + + ``command_id`` comes from the ``StartEvent`` of a prior + :meth:`execute_command_streaming`. Sends a start message, one data message + per chunk of ``inputs``, then (unless ``close_stdin`` is false) a close + message that sends EOF to the process's stdin. Returns the server's + acknowledgement once the input stream is fully sent. + + Not retried: forwarding input is stateful and side-effecting, so a + transport failure is raised to the caller. + + .. py:method:: update_sandbox(name: str, sandbox: Sandbox, update_mask: FieldMask) -> Sandbox Updates mutable fields on an existing Sandbox. Allowlisted update_mask paths today: display_name, diff --git a/pyproject.toml b/pyproject.toml index edcbe610f..31bf19298 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -51,7 +51,8 @@ dev = [ "openai", "langchain-openai", "httpx", - "grpcio<2", # the sandbox extra's dep; also needed to run the gRPC transport unit tests + "grpcio>=1.81.1,<2", # the sandbox extra's dep; also needed to run the gRPC transport unit tests + "build", # some integration tests depend on the databricks-sdk-py wheel "check-manifest", ] notebook = [ @@ -64,15 +65,20 @@ openai = [ "httpx", ] sandbox = [ - # gRPC transport for the sandbox command-execution mixin. Optional so - # users who don't use the sandbox APIs don't pull the large native grpcio - # wheel; imported lazily in databricks/sdk/mixins/_grpc_transport.py. + # gRPC transport and committed stubs for the sandbox command-execution + # mixin. Optional so users who don't use the sandbox APIs don't pull the + # large native grpcio wheel; grpcio and the stubs are imported lazily (see + # databricks/sdk/mixins/_grpc_transport.py and mixins/sandbox.py). # - # `<2` caps the major version to guard against a breaking grpcio 2.x. No - # lower bound is needed: the transport uses only long-stable grpc APIs, and - # pip installs the newest grpcio with wheels for the user's Python (the SDK - # requires >=3.10), so a too-old release is never selected. - "grpcio<2", + # The floors match the version checks baked into the committed stubs: the + # generated sandbox_service_pb2.py validates protobuf >= 6.33.5 and + # sandbox_service_pb2_grpc.py validates grpcio >= 1.81.1 at import, so an + # older runtime that satisfied only the core protobuf floor / an unbounded + # grpcio would raise on import. `<2` caps grpcio's major to guard against a + # breaking 2.x. protobuf's core dependency (above) supplies the `<7` cap and + # the CVE exclusions. + "protobuf>=6.33.5", + "grpcio>=1.81.1,<2", ] [tool.setuptools.dynamic] @@ -87,7 +93,8 @@ include = ["databricks", "databricks.*"] [tool.ruff] line-length = 120 indent-width = 4 -target-version = "py310" +target-version = "py39" +extend-exclude = ["*_pb2.py", "*_pb2_grpc.py"] [tool.ruff.format] quote-style = "double" diff --git a/tests/test_sandbox_mixin.py b/tests/test_sandbox_mixin.py new file mode 100644 index 000000000..fffec0750 --- /dev/null +++ b/tests/test_sandbox_mixin.py @@ -0,0 +1,216 @@ +import pytest + +from databricks.sdk._sandbox_grpc import sandbox_service_pb2 as pb +from databricks.sdk._sandbox_grpc import sandbox_service_pb2_grpc as pb_grpc +from databricks.sdk.mixins import sandbox as sandbox_mod +from databricks.sdk.mixins.sandbox import SandboxExt + +# The fakes below stand in for the authenticated gRPC channel and stub so these +# unit tests exercise execute_command's routing, streaming, retry, and channel +# lifecycle without opening a real network connection. + + +class _FakeApiClient: + def __init__(self): + self._cfg = object() + + +class _FakeChannel: + def __init__(self): + self.closed = 0 + + def close(self): + self.closed += 1 + + +def _response(seq, **event): + return pb.ExecuteCommandResponse(process_event=pb.ProcessEvent(**event), sequence_number=seq) + + +def test_execute_command_streams_responses_and_closes_channel(monkeypatch): + channel = _FakeChannel() + captured = {} + + class _FakeStub: + def __init__(self, ch): + captured["channel"] = ch + + def ExecuteCommand(self, request, metadata=None): + captured["request"] = request + captured["metadata"] = metadata + return iter( + [ + _response(1, start=pb.StartEvent(pid=7)), + _response(2, data=pb.DataEvent(stdout=b"hi\n")), + _response(3, end=pb.EndEvent(exit_code=0)), + ] + ) + + monkeypatch.setattr(sandbox_mod, "open_channel", lambda cfg: channel) + monkeypatch.setattr(pb_grpc, "SandboxServiceStub", _FakeStub) + + api = SandboxExt(_FakeApiClient()) + events = list(api.execute_command_streaming("sandboxes/s1", "/bin/echo", args=["hi"], envs={"K": "v"})) + + assert [e.sequence_number for e in events] == [1, 2, 3] + assert events[0].process_event.start.pid == 7 + assert events[2].process_event.end.exit_code == 0 + assert captured["channel"] is channel + assert captured["request"].cmd == "/bin/echo" + assert list(captured["request"].args) == ["hi"] + assert dict(captured["request"].envs) == {"K": "v"} + assert captured["metadata"] == (("x-databricks-sandbox-name", "sandboxes/s1"),) + assert channel.closed == 1 + + +def test_execute_command_closes_channel_when_abandoned(monkeypatch): + channel = _FakeChannel() + + class _FakeStub: + def __init__(self, ch): + pass + + def ExecuteCommand(self, request, metadata=None): + return iter([_response(1, start=pb.StartEvent(pid=1))]) + + monkeypatch.setattr(sandbox_mod, "open_channel", lambda cfg: channel) + monkeypatch.setattr(pb_grpc, "SandboxServiceStub", _FakeStub) + + stream = SandboxExt(_FakeApiClient()).execute_command_streaming("s", "cmd") + assert next(stream).sequence_number == 1 + stream.close() + assert channel.closed == 1 + + +def test_execute_command_does_not_retry_by_default(monkeypatch): + channel = _FakeChannel() + calls = {"count": 0} + + class _FailingStub: + def __init__(self, ch): + pass + + def ExecuteCommand(self, request, metadata=None): + calls["count"] += 1 + raise RuntimeError("transport failed") + + monkeypatch.setattr(sandbox_mod, "open_channel", lambda cfg: channel) + monkeypatch.setattr(pb_grpc, "SandboxServiceStub", _FailingStub) + + stream = SandboxExt(_FakeApiClient()).execute_command_streaming("s", "cmd") + with pytest.raises(RuntimeError, match="transport failed"): + next(stream) + + assert calls["count"] == 1 + assert channel.closed == 1 + + +def test_execute_command_retries_stream_open_when_opted_in(monkeypatch): + channel = _FakeChannel() + calls = {"count": 0} + + class _FlakyStub: + def __init__(self, ch): + pass + + def ExecuteCommand(self, request, metadata=None): + calls["count"] += 1 + if calls["count"] == 1: + raise RuntimeError("transient") + return iter([_response(1, end=pb.EndEvent(exit_code=0))]) + + monkeypatch.setattr(sandbox_mod, "open_channel", lambda cfg: channel) + monkeypatch.setattr(pb_grpc, "SandboxServiceStub", _FlakyStub) + + events = list( + SandboxExt(_FakeApiClient()).execute_command_streaming( + "s", + "cmd", + is_retryable=lambda _exc: True, + max_attempts=2, + retry_interval_seconds=0.0, + ) + ) + + assert calls["count"] == 2 + assert events[0].process_event.end.exit_code == 0 + assert channel.closed == 1 + + +def test_attach_command_streams_and_closes_channel(monkeypatch): + channel = _FakeChannel() + captured = {} + + class _FakeStub: + def __init__(self, ch): + pass + + def AttachCommand(self, request, metadata=None): + captured["request"] = request + captured["metadata"] = metadata + return iter( + [ + pb.AttachCommandResponse( + process_event=pb.ProcessEvent(end=pb.EndEvent(exit_code=0)), sequence_number=5 + ) + ] + ) + + monkeypatch.setattr(sandbox_mod, "open_channel", lambda cfg: channel) + monkeypatch.setattr(pb_grpc, "SandboxServiceStub", _FakeStub) + + events = list(SandboxExt(_FakeApiClient()).attach_command("sandboxes/s1", "cmd-1", last_sequence_number=4)) + + assert [e.sequence_number for e in events] == [5] + assert captured["request"].command_id == "cmd-1" + assert captured["request"].last_sequence_number == 4 + assert captured["metadata"] == (("x-databricks-sandbox-name", "sandboxes/s1"),) + assert channel.closed == 1 + + +def test_stream_input_sends_start_data_close_and_closes_channel(monkeypatch): + channel = _FakeChannel() + captured = {} + + class _FakeStub: + def __init__(self, ch): + pass + + def StreamInput(self, request_iterator, metadata=None): + captured["requests"] = list(request_iterator) + captured["metadata"] = metadata + return pb.StreamInputResponse() + + monkeypatch.setattr(sandbox_mod, "open_channel", lambda cfg: channel) + monkeypatch.setattr(pb_grpc, "SandboxServiceStub", _FakeStub) + + resp = SandboxExt(_FakeApiClient()).stream_input("sandboxes/s1", "cmd-1", [b"a", b"bc"]) + + assert isinstance(resp, pb.StreamInputResponse) + reqs = captured["requests"] + assert [r.WhichOneof("event") for r in reqs] == ["start", "data", "data", "close"] + assert reqs[0].start.command_id == "cmd-1" + assert [reqs[1].data.data, reqs[2].data.data] == [b"a", b"bc"] + assert captured["metadata"] == (("x-databricks-sandbox-name", "sandboxes/s1"),) + assert channel.closed == 1 + + +def test_stream_input_omits_close_when_disabled(monkeypatch): + channel = _FakeChannel() + captured = {} + + class _FakeStub: + def __init__(self, ch): + pass + + def StreamInput(self, request_iterator, metadata=None): + captured["requests"] = list(request_iterator) + return pb.StreamInputResponse() + + monkeypatch.setattr(sandbox_mod, "open_channel", lambda cfg: channel) + monkeypatch.setattr(pb_grpc, "SandboxServiceStub", _FakeStub) + + SandboxExt(_FakeApiClient()).stream_input("sandboxes/s1", "cmd-1", [b"x"], close_stdin=False) + + assert [r.WhichOneof("event") for r in captured["requests"]] == ["start", "data"] + assert channel.closed == 1