diff --git a/prefect_coiled/credentials.py b/prefect_coiled/credentials.py index 1cdfa6c..c91f54c 100644 --- a/prefect_coiled/credentials.py +++ b/prefect_coiled/credentials.py @@ -7,3 +7,4 @@ class CoiledCredentials(Block): api_token: Optional[SecretStr] = Field(default=None, description="Coiled API token") + workspace: Optional[str] = Field(default=None, description="Coiled Workspace") diff --git a/prefect_coiled/worker.py b/prefect_coiled/worker.py index 83f51f9..cdd6c49 100644 --- a/prefect_coiled/worker.py +++ b/prefect_coiled/worker.py @@ -193,17 +193,24 @@ async def run( else {} ) + workspace = configuration.workspace + # submit the job to run on Coiled creds_config = {} - if configuration.credentials and configuration.credentials.api_token: - creds_config = { - "coiled.token": configuration.credentials.api_token.get_secret_value() - } + if configuration.credentials: + if configuration.credentials.api_token: + creds_config["coiled.token"] = configuration.credentials.api_token.get_secret_value() + + # Workspace for the pool takes precedence (if set), then workspace for credentials. + # If neither is set, then Coiled will use the default workspace set for user in the Coiled web app. + if configuration.credentials.workspace and not configuration.workspace: + workspace = configuration.credentials.workspace + creds_config["coiled.workspace"] = workspace with dask.config.set(creds_config): run_info = run( command=configuration.command, - workspace=configuration.workspace, + workspace=workspace, container=configuration.image if not configuration.software else None, software=configuration.software, secret_env=configuration.env,