reformat

zapata-engineering · May 22, 2024 · ca95b29 · ca95b29
1 parent 7716996
commit ca95b29
Show file tree

Hide file tree

Showing 90 changed files with 219 additions and 212 deletions.
diff --git a/projects/orquestra-sdk/docs/examples/tests/test_dependency_installation.py b/projects/orquestra-sdk/docs/examples/tests/test_dependency_installation.py
@@ -5,9 +5,9 @@
 import typing as t
 
 import pytest
+from orquestra.workflow_shared.schema import ir
 
 from orquestra import sdk
-from orquestra.workflow_shared.schema import ir
 
 
 class Snippets:

diff --git a/projects/orquestra-sdk/docs/examples/tests/test_migrating_from_quantum_engine.py b/projects/orquestra-sdk/docs/examples/tests/test_migrating_from_quantum_engine.py
@@ -5,9 +5,9 @@
 import typing as t
 
 import pytest
+from orquestra.workflow_shared.schema import ir
 
 from orquestra import sdk
-from orquestra.workflow_shared.schema import ir
 
 
 class Snippets:

diff --git a/projects/orquestra-sdk/docs/examples/tests/test_mlflow_utilities.py b/projects/orquestra-sdk/docs/examples/tests/test_mlflow_utilities.py
@@ -14,11 +14,11 @@
 import pytest
 import responses
 import responses.matchers
+from orquestra.workflow_shared.schema import configs
 from pytest import MonkeyPatch
 
 from orquestra import sdk
 from orquestra.sdk._client._base import _config
-from orquestra.workflow_shared.schema import configs
 
 
 class Snippets:

diff --git a/projects/orquestra-sdk/docs/examples/tests/test_secrets.py b/projects/orquestra-sdk/docs/examples/tests/test_secrets.py
@@ -18,10 +18,10 @@
 
 import pytest
 import responses
+from orquestra.workflow_shared.schema import configs
 
 import orquestra.sdk._client._base._config
 import orquestra.sdk._client.secrets._client
-from orquestra.workflow_shared.schema import configs
 
 
 class Snippets:

diff --git a/projects/orquestra-sdk/src/orquestra/sdk/__init__.py b/projects/orquestra-sdk/src/orquestra/sdk/__init__.py
@@ -3,6 +3,10 @@
 ################################################################################
 """Orquestra SDK allows to define computational workflows using Python DSL."""
 
+from orquestra.workflow_shared import Project, ProjectRef, Workspace
+from orquestra.workflow_shared.logs import LogOutput, WorkflowLogs
+from orquestra.workflow_shared.schema.workflow_run import State
+
 from ._client import mlflow, secrets
 from ._client._base._api import (
     CurrentExecutionCtx,
@@ -38,9 +42,6 @@
     WorkflowTemplate,
     workflow,
 )
-from orquestra.workflow_shared import Project, ProjectRef, Workspace
-from orquestra.workflow_shared.logs import LogOutput, WorkflowLogs
-from orquestra.workflow_shared.schema.workflow_run import State
 
 __all__ = [
     "ArtifactFuture",

diff --git a/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/_api/_config.py b/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/_api/_config.py
@@ -7,8 +7,6 @@
 import typing as t
 import warnings
 
-from packaging.version import parse as parse_version
-
 from orquestra.workflow_shared.abc import RuntimeInterface
 from orquestra.workflow_shared.exceptions import (
     ConfigFileNotFoundError,
@@ -20,6 +18,7 @@
     RuntimeConfiguration,
     RuntimeName,
 )
+from packaging.version import parse as parse_version
 
 from .. import _config
 from .._factory import build_runtime_from_config

diff --git a/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/_api/_task_run.py b/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/_api/_task_run.py
@@ -8,7 +8,10 @@
 
 from orquestra.workflow_shared import serde
 from orquestra.workflow_shared.abc import ArtifactValue, RuntimeInterface
-from orquestra.workflow_shared.exceptions import TaskRunNotFound, WorkflowRunIDNotFoundError
+from orquestra.workflow_shared.exceptions import (
+    TaskRunNotFound,
+    WorkflowRunIDNotFoundError,
+)
 from orquestra.workflow_shared.exec_ctx import ExecContext, get_current_exec_context
 from orquestra.workflow_shared.logs import LogOutput
 from orquestra.workflow_shared.schema import ir

diff --git a/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/_api/_wf_run.py b/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/_api/_wf_run.py
@@ -44,7 +44,9 @@
 from orquestra.workflow_shared.schema.workflow_run import ProjectId, State
 from orquestra.workflow_shared.schema.workflow_run import TaskRun as TaskRunModel
 from orquestra.workflow_shared.schema.workflow_run import TaskRunId
-from orquestra.workflow_shared.schema.workflow_run import WorkflowRun as WorkflowRunModel
+from orquestra.workflow_shared.schema.workflow_run import (
+    WorkflowRun as WorkflowRunModel,
+)
 from orquestra.workflow_shared.schema.workflow_run import (
     WorkflowRunId,
     WorkflowRunMinimal,

diff --git a/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/_config.py b/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/_config.py
@@ -12,8 +12,6 @@
 from urllib.parse import ParseResult, urlparse
 
 import filelock
-from pydantic import ValidationError
-
 import orquestra.workflow_shared.exceptions as exceptions
 from orquestra.workflow_shared.schema.configs import (
     CONFIG_FILE_CURRENT_VERSION,
@@ -22,6 +20,7 @@
     RuntimeConfigurationFile,
     RuntimeName,
 )
+from pydantic import ValidationError
 
 from ._env import (
     CONFIG_PATH_ENV,

diff --git a/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/_driver/_client.py b/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/_driver/_client.py
@@ -17,8 +17,6 @@
 
 import pydantic
 import requests
-from requests import codes
-
 from orquestra.workflow_shared import VERSION_REGEX, ProjectRef, exceptions
 from orquestra.workflow_shared.orqdantic import TypeAdapter
 from orquestra.workflow_shared.schema.ir import WorkflowDef
@@ -33,6 +31,7 @@
     WorkflowRunSummary,
     WorkspaceId,
 )
+from requests import codes
 
 from .._spaces._api import make_workspace_zri
 from . import _exceptions, _models

diff --git a/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/_driver/_models.py b/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/_driver/_models.py
@@ -16,8 +16,6 @@
 )
 
 import pydantic
-from typing_extensions import Annotated
-
 from orquestra.workflow_shared.dates import Instant
 from orquestra.workflow_shared.orqdantic import BaseModel
 from orquestra.workflow_shared.schema.ir import WorkflowDef
@@ -31,6 +29,7 @@
     WorkflowRunSummary,
     WorkspaceId,
 )
+from typing_extensions import Annotated
 
 WorkflowDefID = str
 WorkflowRunID = str

diff --git a/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/_dsl.py b/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/_dsl.py
@@ -37,9 +37,6 @@
     from pip_api._parse_requirements import Requirement
 
 import wrapt  # type: ignore
-
-# Needed for fully-qualified type annotations.
-import orquestra.sdk
 from orquestra.workflow_shared.exceptions import (
     DirtyGitRepo,
     InvalidTaskDefinitionError,
@@ -48,6 +45,9 @@
 from orquestra.workflow_shared.kubernetes.quantity import parse_quantity
 from orquestra.workflow_shared.packaging import PackagingError, get_installed_version
 
+# Needed for fully-qualified type annotations.
+import orquestra.sdk
+
 from . import _ast
 
 # ----- SDK exceptions  -----

diff --git a/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/_in_process_runtime.py b/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/_in_process_runtime.py
@@ -8,7 +8,6 @@
 from contextlib import contextmanager
 from datetime import timedelta
 
-from orquestra.sdk import secrets
 from orquestra.workflow_shared import (
     ProjectRef,
     abc,
@@ -32,6 +31,8 @@
     WorkspaceId,
 )
 
+from orquestra.sdk import secrets
+
 WfRunId = str
 ArtifactValue = t.Any
 TaskOutputs = t.Tuple[ArtifactValue, ...]

diff --git a/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/_jwt.py b/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/_jwt.py
@@ -2,7 +2,6 @@
 # © Copyright 2022-2023 Zapata Computing Inc.
 ################################################################################
 import jwt
-
 from orquestra.workflow_shared.exceptions import ExpiredTokenError, InvalidTokenError
 
 

diff --git a/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/_traversal.py b/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/_traversal.py
@@ -15,16 +15,14 @@
 from collections import OrderedDict
 from functools import singledispatch
 
-from pip_api._parse_requirements import Requirement
-
-from orquestra.workflow_shared import exceptions, serde
+from orquestra.workflow_shared import exceptions, parse_git_url, serde
 from orquestra.workflow_shared.exec_ctx import workflow_build
 from orquestra.workflow_shared.packaging import (
     get_current_python_version,
     get_current_sdk_version,
 )
 from orquestra.workflow_shared.schema import ir, responses
-from orquestra.workflow_shared import parse_git_url
+from pip_api._parse_requirements import Requirement
 
 from . import _dsl, _workflow
 

diff --git a/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/_viz.py b/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/_viz.py
@@ -8,7 +8,6 @@
 from dataclasses import dataclass
 
 import graphviz  # type: ignore
-
 from orquestra.workflow_shared.schema import ir
 
 NodeId = str

diff --git a/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/_workflow.py b/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/_workflow.py
@@ -22,13 +22,13 @@
     overload,
 )
 
-from typing_extensions import ParamSpec
-
 import orquestra.workflow_shared.schema.ir as ir
-from orquestra.sdk import secrets
 from orquestra.workflow_shared import exceptions
 from orquestra.workflow_shared.exceptions import WorkflowSyntaxError
 from orquestra.workflow_shared.schema.workflow_run import ProjectId, WorkspaceId
+from typing_extensions import ParamSpec
+
+from orquestra.sdk import secrets
 
 from . import _api, _dsl, loader
 from ._ast import CallVisitor, NodeReference, NodeReferenceType, normalize_indents

diff --git a/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/cli/_dumpers.py b/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/cli/_dumpers.py
@@ -9,7 +9,10 @@
 
 from orquestra.workflow_shared import serde
 from orquestra.workflow_shared.logs import LogOutput, WorkflowLogs
-from orquestra.workflow_shared.schema.workflow_run import TaskInvocationId, WorkflowRunId
+from orquestra.workflow_shared.schema.workflow_run import (
+    TaskInvocationId,
+    WorkflowRunId,
+)
 
 
 class WFOutputDumper:

diff --git a/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/cli/_entry.py b/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/cli/_entry.py
@@ -10,7 +10,6 @@
 
 import click
 import cloup
-
 from orquestra.workflow_shared.schema.configs import RuntimeName
 
 from . import _cli_logs

diff --git a/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/cli/_login/_login.py b/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/cli/_login/_login.py
@@ -6,8 +6,10 @@
 import typing as t
 
 from aiohttp import web
-
-from orquestra.workflow_shared.exceptions import LocalConfigLoginError, UserCancelledPrompt
+from orquestra.workflow_shared.exceptions import (
+    LocalConfigLoginError,
+    UserCancelledPrompt,
+)
 from orquestra.workflow_shared.schema.configs import RemoteRuntime
 
 from .. import _arg_resolvers, _repos

diff --git a/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/cli/_repos.py b/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/cli/_repos.py
@@ -15,9 +15,6 @@
 from types import ModuleType
 
 import requests
-from typing_extensions import assert_never
-
-from orquestra import sdk
 from orquestra.workflow_shared import exceptions
 from orquestra.workflow_shared.abc import ArtifactValue
 from orquestra.workflow_shared.dates import from_unix_time
@@ -40,6 +37,9 @@
     WorkflowRunSummary,
     WorkspaceId,
 )
+from typing_extensions import assert_never
+
+from orquestra import sdk
 
 from ..._base import _config, loader
 from ..._base._driver._client import DriverClient, ExternalUriProvider

diff --git a/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/cli/_task/_logs.py b/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/cli/_task/_logs.py
@@ -6,7 +6,10 @@
 from pathlib import Path
 
 from orquestra.workflow_shared.schema.configs import ConfigName
-from orquestra.workflow_shared.schema.workflow_run import TaskInvocationId, WorkflowRunId
+from orquestra.workflow_shared.schema.workflow_run import (
+    TaskInvocationId,
+    WorkflowRunId,
+)
 
 from .. import _arg_resolvers, _dumpers, _repos
 from .._ui import _presenters

diff --git a/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/cli/_ui/_errors.py b/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/cli/_ui/_errors.py
@@ -9,13 +9,12 @@
 
 import click
 import rich
+from orquestra.workflow_shared import exceptions
+from orquestra.workflow_shared.schema.responses import ResponseStatusCode
 from rich.box import SIMPLE_HEAVY
 from rich.console import Console
 from rich.table import Column, Table
 
-from orquestra.workflow_shared import exceptions
-from orquestra.workflow_shared.schema.responses import ResponseStatusCode
-
 from ...._base._config import IN_PROCESS_CONFIG_NAME, RAY_CONFIG_NAME_ALIAS
 
 

diff --git a/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/cli/_ui/_presenters.py b/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/cli/_ui/_presenters.py
@@ -15,15 +15,6 @@
 from typing import List, Optional, Sequence
 
 import click
-from rich.box import SIMPLE_HEAVY
-from rich.console import Console, Group, RenderableType
-from rich.live import Live
-from rich.pretty import Pretty
-from rich.rule import Rule
-from rich.spinner import Spinner
-from rich.table import Column, Table
-from tabulate import tabulate
-
 from orquestra.workflow_shared import serde
 from orquestra.workflow_shared.dates import Instant, from_unix_time
 from orquestra.workflow_shared.logs import LogOutput, WorkflowLogs
@@ -34,7 +25,18 @@
     RuntimeName,
 )
 from orquestra.workflow_shared.schema.ir import ArtifactFormat
-from orquestra.workflow_shared.schema.workflow_run import TaskInvocationId, WorkflowRunId
+from orquestra.workflow_shared.schema.workflow_run import (
+    TaskInvocationId,
+    WorkflowRunId,
+)
+from rich.box import SIMPLE_HEAVY
+from rich.console import Console, Group, RenderableType
+from rich.live import Live
+from rich.pretty import Pretty
+from rich.rule import Rule
+from rich.spinner import Spinner
+from rich.table import Column, Table
+from tabulate import tabulate
 
 from ...._base import _env
 from . import _errors

diff --git a/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/cli/_ui/_prompts.py b/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/cli/_ui/_prompts.py
@@ -7,7 +7,6 @@
 from typing import overload
 
 from inquirer import errors  # type: ignore # noqa
-
 from orquestra.workflow_shared import exceptions
 
 # One of our transitive dependencies shows DeprecationWarnings related to invalid usage

diff --git a/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/cli/_workflow/_list.py b/projects/orquestra-sdk/src/orquestra/sdk/_client/_base/cli/_workflow/_list.py
@@ -6,7 +6,10 @@
 
 from orquestra.workflow_shared import exceptions as exceptions
 from orquestra.workflow_shared.schema.configs import ConfigName
-from orquestra.workflow_shared.schema.workflow_run import WorkflowRunSummary, WorkspaceId
+from orquestra.workflow_shared.schema.workflow_run import (
+    WorkflowRunSummary,
+    WorkspaceId,
+)
 
 from .. import _arg_resolvers, _repos
 from .._ui import _presenters

diff --git a/projects/orquestra-sdk/src/orquestra/sdk/_client/mlflow/_connection_utils.py b/projects/orquestra-sdk/src/orquestra/sdk/_client/mlflow/_connection_utils.py
@@ -10,10 +10,13 @@
 from pathlib import Path
 from typing import Optional, Tuple
 
+from orquestra.workflow_shared.exceptions import (
+    ConfigNameNotFoundError,
+    RuntimeConfigError,
+)
 from requests import Response, Session
 
 from orquestra import sdk
-from orquestra.workflow_shared.exceptions import ConfigNameNotFoundError, RuntimeConfigError
 
 from .._base import _env
 from .._base._config import read_config