[ISSUE] ImportError: cannot import name 'ServedModelInputWorkloadSize' from 'databricks.sdk.service.serving'

[scardonal]

Description
I'm having this error with Databricks Runtime 16.4 LTS and databricks-sdk==0.55.0

ImportError: cannot import name 'ServedModelInputWorkloadSize' from 'databricks.sdk.service.serving'

Because of that, I can not use databricks-agents library

Reproduction
from databricks import agents

Expected behavior
A clear and concise description of what you expected to happen.

Is it a regression?
Did this work in a previous version of the SDK? If so, which versions did you try?
Yes, it did, at least with version 0.50.0

Debug Logs

ImportError: cannot import name 'ServedModelInputWorkloadSize' from 'databricks.sdk.service.serving' (/local_disk0/.ephemeral_nfs/envs/pythonEnv-4ad4068a-f11c-45ed-a3e0-419f392fde4a/lib/python3.12/site-packages/databricks/sdk/service/serving.py)
File <command-7073579183174803>, line 1
----> 1 from databricks import agents
File /databricks/python_shell/lib/dbruntime/autoreload/discoverability/hook.py:71, in AutoreloadDiscoverabilityHook._patched_import(self, name, *args, **kwargs)
     65 if not self._should_hint and (
     66     (module := sys.modules.get(absolute_name)) is not None and
     67     (fname := get_allowed_file_name_or_none(module)) is not None and
     68     (mtime := os.stat(fname).st_mtime) > self.last_mtime_by_modname.get(
     69         absolute_name, float("inf")) and not self._should_hint):
     70     self._should_hint = True
---> 71 module = self._original_builtins_import(name, *args, **kwargs)
     72 if (fname := fname or get_allowed_file_name_or_none(module)) is not None:
     73     mtime = mtime or os.stat(fname).st_mtime
File /local_disk0/.ephemeral_nfs/envs/pythonEnv-4ad4068a-f11c-45ed-a3e0-419f392fde4a/lib/python3.12/site-packages/databricks/agents/__init__.py:1
----> 1 from databricks.agents.deployments import (
      2     delete_deployment,
      3     deploy,
      4     get_deployments,
      5     list_deployments,
      6 )
      7 from databricks.agents.permissions import get_permissions, set_permissions
      8 from databricks.agents.reviews import (
      9     enable_trace_reviews,
     10     get_review_instructions,
     11     set_review_instructions,
     12 )
File /local_disk0/.ephemeral_nfs/envs/pythonEnv-4ad4068a-f11c-45ed-a3e0-419f392fde4a/lib/python3.12/site-packages/databricks/agents/deployments.py:53
     45 from databricks.sdk import WorkspaceClient
     46 from databricks.sdk.errors.platform import (
     47     BadRequest,
     48     InvalidParameterValue,
   (...)
     51     ResourceDoesNotExist,
     52 )
---> 53 from databricks.sdk.service.serving import (
     54     AiGatewayConfig,
     55     AiGatewayInferenceTableConfig,
     56     AutoCaptureConfigInput,
     57     EndpointCoreConfigInput,
     58     EndpointCoreConfigOutput,
     59     EndpointPendingConfig,
     60     EndpointTag,
     61     Route,
     62     ServedEntityInput,
     63     ServedEntityOutput,
     64     ServedModelInputWorkloadSize,
     65     ServingEndpointDetailed,
     66     TrafficConfig,
     67 )
     69 _logger = logging.getLogger("agents")
     71 __DEPLOY_ENV_VARS_WITH_STATIC_VALUES = {
     72     "ENABLE_LANGCHAIN_STREAMING": "true",
     73     "ENABLE_MLFLOW_TRACING": "true",
     74     "RETURN_REQUEST_ID_IN_RESPONSE": "true",
     75 }

Other Information

• OS: [e.g. macOS] Databricks Runtime 16.4 LTS
• Version: [e.g. 0.1.0] 0.55.0

Additional context

[veenaramesh]

Hey @scardonal-- what version of databricks-agents are you using? In the latest version, agents.deploy() takes workload_size as a string (docs here). So, you should not need to import ServedModelInputWorkloadSize at all.

[scardonal]

Hey @veenaramesh . I'm not sure now because I've already moved on with my project.

When I faced this error, I couldn't even import the agents module. I remember that I was using the default version of databricks-agents that comes with the runtime 16.4 LTS