move secrets and fix tests

projects/orquestra-sdk/src/orquestra/sdk/__init__.py

@@ -3,10 +3,10 @@
 ################################################################################
 """Orquestra SDK allows to define computational workflows using Python DSL."""
 
-from orquestra.workflow_runtime import secrets
-from orquestra.workflow_shared import Project, ProjectRef, Workspace
+from orquestra.workflow_shared import Project, ProjectRef, Workspace, secrets
 from orquestra.workflow_shared.logs import LogOutput, WorkflowLogs
 from orquestra.workflow_shared.schema.workflow_run import State
+from orquestra.workflow_shared.secrets import Secret
 
 from ._client import mlflow
 from ._client._base._api import (
@@ -32,7 +32,6 @@
     LocalImport,
     PythonImports,
     Resources,
-    Secret,
     TaskDef,
     task,
 )

projects/orquestra-sdk/src/orquestra/sdk/_client/_base/_dsl.py

@@ -44,6 +44,7 @@
 )
 from orquestra.workflow_shared.kubernetes.quantity import parse_quantity
 from orquestra.workflow_shared.packaging import PackagingError, get_installed_version
+from orquestra.workflow_shared.secrets import Secret
 
 # Needed for fully-qualified type annotations.
 import orquestra.sdk
@@ -78,50 +79,9 @@ class UnknownPlaceholderInCustomNameWarning(Warning):
 # function calls).
 Argument = Union[Constant, "ArtifactFuture", "Secret"]
 
-_secret_as_string_error = (
-    "Invalid usage of a Secret object. Secrets are not "
-    "available when building the workflow graph and cannot"
-    " be used as strings. If you need to use a Secret's"
-    " value, this must be done inside of a task."
-)
-
-
 _TaskReturn = TypeVar("_TaskReturn")
 
 
-class Secret(NamedTuple):
-    name: str
-    # Config name is only used for the local runtimes where we can't infer the location
-    # where we get a secret's value from.
-    # This matches the behaviour of `sdk.secrets.get` where the config name is used to
-    # get a secret when running locally.
-    config_name: Optional[str] = None
-    # Workspace ID is used by local and remote runtimes to fetch a secret from a
-    # specific workspace.
-    workspace_id: Optional[str] = None
-
-    def __reduce__(self) -> str | tuple[Any, ...]:
-        # We need to override the pickling behaviour for Secret
-        # This is because we override other dunder methods which cause the normal
-        # picling behaviour to fail.
-        return (self.__class__, (self.name, self.config_name, self.workspace_id))
-
-    def __getattr__(self, item):
-        try:
-            return self.__getattribute__(item)
-        except AttributeError as e:
-            raise AttributeError(_secret_as_string_error) from e
-
-    def __getitem__(self, item):
-        raise AttributeError(_secret_as_string_error)
-
-    def __str__(self):
-        raise AttributeError(_secret_as_string_error)
-
-    def __iter__(self):
-        raise AttributeError(_secret_as_string_error)
-
-
 @dataclass(frozen=True, eq=True)
 class GitImportWithAuth:
     """A task import that uses a private Git repo.

projects/orquestra-sdk/src/orquestra/sdk/_client/_base/_traversal.py

@@ -22,6 +22,7 @@
     get_current_sdk_version,
 )
 from orquestra.workflow_shared.schema import ir, responses
+from orquestra.workflow_shared.secrets import Secret
 from pip_api._parse_requirements import Requirement
 
 from . import _dsl, _workflow
@@ -254,7 +255,7 @@ def traverse(self, output_nodes: t.Sequence[DSLDataNode]):
                 seen_futures.add(n)
                 seen_invocations.add(n.invocation)
 
-            elif isinstance(n, _dsl.Secret):
+            elif isinstance(n, Secret):
                 self._secrets[_make_key(n)] = ir.SecretNode(
                     id=f"secret-{secret_counter}",
                     secret_name=n.name,
@@ -315,7 +316,7 @@ def get_node_id(self, node: DSLDataNode) -> ir.ArgumentId:
         if isinstance(node, _dsl.ArtifactFuture):
             return self._future_artifacts[node].id
 
-        elif isinstance(node, _dsl.Secret):
+        elif isinstance(node, Secret):
             return self._secrets[key].id
         else:
             return self._constants[key].id

projects/orquestra-sdk/src/orquestra/sdk/_client/_base/_workflow.py

@@ -26,6 +26,7 @@
 from orquestra.workflow_shared import exceptions
 from orquestra.workflow_shared.exceptions import WorkflowSyntaxError
 from orquestra.workflow_shared.schema.workflow_run import ProjectId, WorkspaceId
+from orquestra.workflow_shared.secrets import Secret
 from typing_extensions import ParamSpec
 
 from orquestra.sdk import secrets
@@ -37,7 +38,6 @@
     FunctionRef,
     Import,
     ImportTypes,
-    Secret,
     TaskDef,
     UnknownPlaceholderInCustomNameWarning,
     get_fn_ref,

projects/orquestra-workflow-runtime/src/orquestra/workflow_runtime/_ray/_build_workflow.py

@@ -14,7 +14,7 @@
 from packaging import version
 from typing_extensions import assert_never
 
-from orquestra.sdk import secrets
+from orquestra.workflow_shared import secrets
 from orquestra.workflow_shared import (
     SEMVER_REGEX,
     dispatch,

projects/orquestra-workflow-runtime/src/orquestra/workflow_runtime/secrets/__init__.py → projects/orquestra-workflow-shared/src/orquestra/workflow_shared/secrets/__init__.py

@@ -19,11 +19,12 @@
 Platform, it uses a server-side authorization mechanism handled automatically.
 """
 
-from ._api import delete, get, list, set
+from ._api import delete, get, list, set, Secret
 
 __all__ = [
     "delete",
     "get",
     "list",
     "set",
+    "Secret",
 ]

projects/orquestra-workflow-runtime/src/orquestra/workflow_runtime/secrets/_api.py → projects/orquestra-workflow-shared/src/orquestra/workflow_shared/secrets/_api.py

@@ -3,20 +3,59 @@
 ################################################################################
 """Code for user-facing utilities related to secrets."""
 import typing as t
+from typing import NamedTuple
 
 from orquestra.workflow_shared import exceptions as sdk_exc
 from orquestra.workflow_shared.exec_ctx import ExecContext, get_current_exec_context
 from orquestra.workflow_shared.schema.configs import ConfigName
 from orquestra.workflow_shared.schema.workflow_run import WorkspaceId
 
-from .._base import _dsl
 from . import _auth, _exceptions, _models
 
 
 def _translate_to_zri(workspace_id: WorkspaceId, secret_name: str) -> str:
     """Create ZRI from workspace_id and secret_name."""
     return f"zri:v1::0:{workspace_id}:secret:{secret_name}"
 
+_secret_as_string_error = (
+    "Invalid usage of a Secret object. Secrets are not "
+    "available when building the workflow graph and cannot"
+    " be used as strings. If you need to use a Secret's"
+    " value, this must be done inside of a task."
+)
+
+class Secret(NamedTuple):
+    name: str
+    # Config name is only used for the local runtimes where we can't infer the location
+    # where we get a secret's value from.
+    # This matches the behaviour of `sdk.secrets.get` where the config name is used to
+    # get a secret when running locally.
+    config_name: t.Optional[str] = None
+    # Workspace ID is used by local and remote runtimes to fetch a secret from a
+    # specific workspace.
+    workspace_id: t.Optional[str] = None
+
+    def __reduce__(self) -> str | tuple[t.Any, ...]:
+        # We need to override the pickling behaviour for Secret
+        # This is because we override other dunder methods which cause the normal
+        # picling behaviour to fail.
+        return (self.__class__, (self.name, self.config_name, self.workspace_id))
+
+    def __getattr__(self, item):
+        try:
+            return self.__getattribute__(item)
+        except AttributeError as e:
+            raise AttributeError(_secret_as_string_error) from e
+
+    def __getitem__(self, item):
+        raise AttributeError(_secret_as_string_error)
+
+    def __str__(self):
+        raise AttributeError(_secret_as_string_error)
+
+    def __iter__(self):
+        raise AttributeError(_secret_as_string_error)
+
 
 def get(
     name: str,
@@ -56,7 +95,7 @@ def get(
     if get_current_exec_context() == ExecContext.WORKFLOW_BUILD:
         return t.cast(
             str,
-            _dsl.Secret(name=name, config_name=config_name, workspace_id=workspace_id),
+            Secret(name=name, config_name=config_name, workspace_id=workspace_id),
         )
 
     try:

projects/orquestra-workflow-runtime/src/orquestra/workflow_runtime/secrets/_auth.py → projects/orquestra-workflow-shared/src/orquestra/workflow_shared/secrets/_auth.py

@@ -8,8 +8,6 @@
 from orquestra.workflow_shared import exceptions
 from orquestra.workflow_shared.schema.configs import ConfigName
 
-from .._base import _config
-from .._base._env import PASSPORT_FILE_ENV
 from ._client import SecretsClient
 
 # We assume that we can access the Config Service under a well-known URI if the passport
@@ -18,6 +16,7 @@
 
 
 def _authorize_with_passport() -> t.Optional[SecretsClient]:
+    from orquestra.sdk._client._base._env import PASSPORT_FILE_ENV
     if (passport_path := os.getenv(PASSPORT_FILE_ENV)) is None:
         return None
 
@@ -26,6 +25,11 @@ def _authorize_with_passport() -> t.Optional[SecretsClient]:
 
 
 def _read_config_opts(config_name: ConfigName):
+    # This import is awful, as it makes shared code dependent on client.
+    # In reality, this function in only called on client side, so from functional
+    # perspective this will work.
+    from orquestra.sdk._client._base import _config
+
     try:
         cfg = _config.read_config(config_name=config_name)
     except exceptions.ConfigNameNotFoundError: