dagster-cloud 1.8.2__py3-none-any.whl → 1.12.6__py3-none-any.whl

dagster_cloud/pex/grpc/types.py

@@ -1,4 +1,4 @@
-from typing import List, NamedTuple, Optional
+from typing import NamedTuple, Optional
 
 import dagster._check as check
 from dagster._core.instance.ref import InstanceRef
@@ -21,7 +21,7 @@ class PexServerHandle(
     )
 ):
     def __new__(cls, deployment_name: str, location_name: str, metadata_update_timestamp: int):
-        return super(PexServerHandle, cls).__new__(
+        return super().__new__(
             cls,
             check.str_param(deployment_name, "deployment_name"),
             check.str_param(location_name, "location_name"),
@@ -49,7 +49,7 @@ class CreatePexServerArgs(
         code_location_deploy_data: CodeLocationDeployData,
         instance_ref: Optional[InstanceRef] = None,
     ):
-        return super(CreatePexServerArgs, cls).__new__(
+        return super().__new__(
             cls,
             check.inst_param(server_handle, "server_handle", PexServerHandle),
             check.inst_param(
@@ -100,7 +100,7 @@ class GetPexServersResponse(
     NamedTuple(
         "_GetPexServersResponse",
         [
-            ("server_handles", List[PexServerHandle]),
+            ("server_handles", list[PexServerHandle]),
         ],
     )
 ):
@@ -112,7 +112,7 @@ class GetCrashedPexServersResponse(
     NamedTuple(
         "_GetCrashedPexServersResponse",
         [
-            ("server_handles", List[PexServerHandle]),
+            ("server_handles", list[PexServerHandle]),
         ],
     )
 ):

dagster_cloud/secrets/__init__.py

@@ -1 +1 @@
-from .loader import DagsterCloudSecretsLoader as DagsterCloudSecretsLoader
+from dagster_cloud.secrets.loader import DagsterCloudSecretsLoader as DagsterCloudSecretsLoader

dagster_cloud/secrets/loader.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Optional
+from typing import Any, Optional
 
 from dagster._core.secrets import SecretsLoader
 from dagster._serdes import ConfigurableClass
@@ -24,9 +24,9 @@ class DagsterCloudSecretsLoader(SecretsLoader, ConfigurableClass):
         self._inst_data = inst_data
 
     def _execute_query(self, query, variables=None):
-        return self._instance.graphql_client.execute(query, variable_values=variables)
+        return self._instance.graphql_client.execute(query, variable_values=variables)  # pyright: ignore[reportAttributeAccessIssue]
 
-    def get_secrets_for_environment(self, location_name: Optional[str]) -> Dict[str, str]:
+    def get_secrets_for_environment(self, location_name: Optional[str]) -> dict[str, str]:
         res = self._execute_query(
             SECRETS_QUERY,
             variables={"locationName": location_name},

dagster_cloud/serverless/__init__.py

@@ -1 +1 @@
-from .io_manager import serverless_io_manager as serverless_io_manager
+from dagster_cloud.serverless.io_manager import serverless_io_manager as serverless_io_manager

dagster_cloud/serverless/io_manager.py

@@ -2,19 +2,20 @@ import datetime
 import io
 import os
 import pickle
-from typing import Sequence, Tuple, Union
+from typing import Any, Union
 
 import boto3
 import dagster._check as check
 import requests
-from dagster import InputContext, IOManager, OutputContext, io_manager
+from dagster import InputContext, OutputContext, UPathIOManager, io_manager
 from dagster._utils import PICKLE_PROTOCOL
 from dagster._vendored.dateutil import parser
+from upath import UPath
 
 ECS_AGENT_IP = "169.254.170.2"
 
 
-class PickledObjectServerlessIOManager(IOManager):
+class PickledObjectServerlessIOManager(UPathIOManager):
     def __init__(
         self,
         s3_bucket,
@@ -23,8 +24,10 @@ class PickledObjectServerlessIOManager(IOManager):
         self._bucket = check.str_param(s3_bucket, "s3_bucket")
         self._s3_prefix = check.str_param(s3_prefix, "s3_prefix")
         self._boto_session, self._boto_session_expiration = self._refresh_boto_session()
+        base_path = UPath(s3_prefix) if s3_prefix else None
+        super().__init__(base_path=base_path)
 
-    def _refresh_boto_session(self) -> Tuple[boto3.Session, datetime.datetime]:
+    def _refresh_boto_session(self) -> tuple[boto3.Session, datetime.datetime]:
         # We have to do this whacky way to get credentials to ensure that we get iam role
         # we assigned to the task. If we used the default boto behavior, it could get overriden
         # when users set AWS env vars.
@@ -36,7 +39,7 @@ class PickledObjectServerlessIOManager(IOManager):
             aws_session_token=aws_creds["Token"],
         )
         expiration = parser.parse(aws_creds["Expiration"])
-        return session, expiration
+        return session, expiration  # pyright: ignore[reportReturnType]
 
     @property
     def _s3(self):
@@ -44,63 +47,43 @@ class PickledObjectServerlessIOManager(IOManager):
             self._boto_session_expiration.tzinfo
         ) + datetime.timedelta(minutes=5):
             self._boto_session, self._boto_session_expiration = self._refresh_boto_session()
-        return self._boto_session.client("s3", region_name="us-west-2")
-
-    def _get_path(self, context: Union[InputContext, OutputContext]) -> str:
-        path: Sequence[str]
-        if context.has_asset_key:
-            path = context.get_asset_identifier()
-        else:
-            path = ["storage", *context.get_identifier()]
-
-        return "/".join([self._s3_prefix, *path])
-
-    def has_output(self, context):
-        key = self._get_path(context)
-        return self._has_object(key)
-
-    def _has_object(self, key):
-        check.str_param(key, "key")
-        check.param_invariant(len(key) > 0, "key")
-
-        found_object = False
+        return self._boto_session.client(
+            "s3",
+            region_name=os.getenv("DAGSTER_CLOUD_SERVERLESS_REGION", "us-west-2"),
+        )
 
+    def load_from_path(self, context: InputContext, path: UPath) -> Any:
         try:
-            self._s3.get_object(Bucket=self._bucket, Key=key)
-            found_object = True
+            s3_obj = self._s3.get_object(Bucket=self._bucket, Key=path.as_posix())["Body"].read()
+            return pickle.loads(s3_obj)
         except self._s3.exceptions.NoSuchKey:
-            found_object = False
-
-        return found_object
+            raise FileNotFoundError(
+                f"Could not find the input for '{context.name}'. It may have expired."
+            )
 
-    def load_input(self, context):
-        if context.dagster_type.typing_type == type(None):
-            return None
+    def dump_to_path(self, context: OutputContext, obj: Any, path: UPath) -> None:
+        pickled_obj = pickle.dumps(obj, PICKLE_PROTOCOL)
+        pickled_obj_bytes = io.BytesIO(pickled_obj)
+        self._s3.upload_fileobj(pickled_obj_bytes, self._bucket, path.as_posix())
 
-        key = self._get_path(context)
-        check.invariant(self._has_object(key), "Input not found. It may have expired.")
-        obj = pickle.loads(self._s3.get_object(Bucket=self._bucket, Key=key)["Body"].read())
+    def path_exists(self, path: UPath) -> bool:
+        try:
+            self._s3.get_object(Bucket=self._bucket, Key=path.as_posix())
+        except self._s3.exceptions.NoSuchKey:
+            return False
+        return True
 
-        return obj
+    def unlink(self, path: UPath) -> None:
+        self._s3.delete_object(Bucket=self._bucket, Key=path.as_posix())
 
-    def handle_output(self, context, obj):
-        if context.dagster_type.typing_type == type(None):
-            check.invariant(
-                obj is None,
-                "Output had Nothing type or 'None' annotation, but handle_output received"
-                f" value that was not None and was of type {type(obj)}.",
-            )
-            return None
+    def make_directory(self, path: UPath) -> None:
+        # It is not necessary to create directories in S3
+        return None
 
-        key = self._get_path(context)
+    def get_op_output_relative_path(self, context: Union[InputContext, OutputContext]):
+        from upath import UPath
 
-        pickled_obj = pickle.dumps(obj, PICKLE_PROTOCOL)
-        pickled_obj_bytes = io.BytesIO(pickled_obj)
-        self._s3.upload_fileobj(
-            pickled_obj_bytes,
-            self._bucket,
-            key,
-        )
+        return UPath(*["storage", *context.get_identifier()])
 
 
 @io_manager

dagster_cloud/storage/client.py

@@ -19,25 +19,62 @@ from dagster_cloud_cli.core.graphql_client import (
 
 def dagster_cloud_api_config():
     return {
-        "url": Field(StringSource, is_required=False),
-        "agent_token": Field(StringSource, is_required=True),
-        "headers": Field(Permissive(), default_value={}),
-        "cookies": Field(Permissive(), default_value={}),
-        "timeout": Field(Noneable(IntSource), default_value=DEFAULT_TIMEOUT),
-        "verify": Field(bool, default_value=True),
-        "retries": Field(IntSource, default_value=DEFAULT_RETRIES),
-        "method": Field(StringSource, default_value="POST"),
-        "backoff_factor": Field(float, default_value=DEFAULT_BACKOFF_FACTOR),
-        "proxies": Field(Map(str, str), is_required=False),
-        "socket_options": Field(Noneable(Array(Array(Any)))),
-        "agent_label": Field(StringSource, is_required=False),
-        # Handle requests for a single non-branch deployment
+        "url": Field(
+            StringSource,
+            is_required=False,
+            description="Dagster+ API server URL. This can be omitted and is derived from your agent token.",
+        ),
+        "agent_token": Field(
+            StringSource, is_required=True, description="Dagster+ agent api token."
+        ),
+        "agent_label": Field(
+            StringSource,
+            is_required=False,
+            description="Custom label visible in the Dagster+ UI.",
+        ),
         "deployment": Field(
-            ScalarUnion(scalar_type=str, non_scalar_schema=Array(str)), is_required=False
+            ScalarUnion(scalar_type=str, non_scalar_schema=Array(str)),
+            is_required=False,
+            description="Handle requests for a single deployment.",
         ),
         # Handle requests for multiple non-branch deployments
-        "deployments": Field(Array(StringSource), is_required=False),
+        "deployments": Field(
+            Array(StringSource),
+            is_required=False,
+            description="Handle requests for multiple deployments",
+        ),
         # Handle requests for all branch deployments (can be combined with `deployment`)`
-        "branch_deployments": Field(BoolSource, default_value=False, is_required=False),
-        "all_serverless_deployments": Field(BoolSource, default_value=False, is_required=False),
+        "branch_deployments": Field(
+            BoolSource,
+            default_value=False,
+            is_required=False,
+            description="Handle requests for all branch deployments (can be combined with `deployment` or `deployments`)",
+        ),
+        "timeout": Field(
+            Noneable(IntSource),
+            default_value=DEFAULT_TIMEOUT,
+            description="How long before a request times out against the Dagster+ API servers.",
+        ),
+        "retries": Field(
+            IntSource,
+            default_value=DEFAULT_RETRIES,
+            description="How many times to retry retriable response codes (429, 503, etc.) before failing.",
+        ),
+        "backoff_factor": Field(
+            float,
+            default_value=DEFAULT_BACKOFF_FACTOR,
+            description="Exponential factor to back off between retries.",
+        ),
+        "method": Field(StringSource, default_value="POST", description="Default HTTP method."),
+        "verify": Field(bool, default_value=True, description="If False, ignore verifying SSL."),
+        "headers": Field(Permissive(), default_value={}, description="Custom headers."),
+        "cookies": Field(Permissive(), default_value={}, description="Custom cookies."),
+        "proxies": Field(Map(str, str), is_required=False, description="Custom proxies."),
+        "socket_options": Field(Noneable(Array(Array(Any))), description="Custom socket options."),
+        "all_serverless_deployments": Field(
+            BoolSource,
+            default_value=False,
+            is_required=False,
+            description="Unused by hybrid agents",
+        ),
     }

dagster_cloud/storage/compute_logs/__init__.py

@@ -1 +1,3 @@
-from .compute_log_manager import CloudComputeLogManager as CloudComputeLogManager
+from dagster_cloud.storage.compute_logs.compute_log_manager import (
+    CloudComputeLogManager as CloudComputeLogManager,
+)

dagster_cloud/storage/compute_logs/compute_log_manager.py

@@ -1,6 +1,7 @@
-from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence
+import os
+from collections.abc import Sequence
+from typing import IO, TYPE_CHECKING, Any, Optional
 
-import dagster._seven as _seven
 import requests
 from dagster import (
     Field,
@@ -8,16 +9,18 @@ from dagster import (
     StringSource,
     _check as check,
 )
-from dagster._core.storage.cloud_storage_compute_log_manager import CloudStorageComputeLogManager
+from dagster._core.storage.cloud_storage_compute_log_manager import (
+    TruncatingCloudStorageComputeLogManager,
+)
 from dagster._core.storage.compute_log_manager import ComputeIOType
 from dagster._core.storage.local_compute_log_manager import (
     IO_TYPE_EXTENSION,
     LocalComputeLogManager,
 )
 from dagster._serdes import ConfigurableClass, ConfigurableClassData
-from dagster._utils import ensure_file
 from dagster_cloud_cli.core.errors import raise_http_error
 from dagster_cloud_cli.core.headers.auth import DagsterCloudInstanceScope
+from dagster_shared import seven
 from requests.adapters import HTTPAdapter
 from typing_extensions import Self
 
@@ -26,7 +29,7 @@ if TYPE_CHECKING:
 
 
 class CloudComputeLogManager(
-    CloudStorageComputeLogManager["DagsterCloudAgentInstance"], ConfigurableClass
+    TruncatingCloudStorageComputeLogManager["DagsterCloudAgentInstance"], ConfigurableClass
 ):
     def __init__(
         self,
@@ -36,7 +39,7 @@ class CloudComputeLogManager(
     ):
         # proxy calls to local compute log manager (for subscriptions, etc)
         if not local_dir:
-            local_dir = _seven.get_system_temp_directory()
+            local_dir = seven.get_system_temp_directory()
 
         self._upload_session = requests.Session()
         adapter = HTTPAdapter(max_retries=3)
@@ -46,6 +49,7 @@ class CloudComputeLogManager(
         self._local_manager = LocalComputeLogManager(local_dir)
         self._upload_interval = check.opt_int_param(upload_interval, "upload_interval")
         self._inst_data = check.opt_inst_param(inst_data, "inst_data", ConfigurableClassData)
+        super().__init__()
 
     @property
     def inst_data(self):
@@ -87,21 +91,23 @@ class CloudComputeLogManager(
         """Returns whether the cloud storage contains logs for a given log key."""
         return False
 
-    def upload_to_cloud_storage(
-        self, log_key: Sequence[str], io_type: ComputeIOType, partial=False
+    def _upload_file_obj(
+        self, data: IO[bytes], log_key: Sequence[str], io_type: ComputeIOType, partial=False
     ):
         path = self.local_manager.get_captured_local_path(log_key, IO_TYPE_EXTENSION[io_type])
-        ensure_file(path)
-        params: Dict[str, Any] = {
+        size_bytes = os.stat(path).st_size
+        params: dict[str, Any] = {
             "log_key": log_key,
             "io_type": io_type.value,
+            "size_bytes": size_bytes,
             # for back-compat
             "run_id": log_key[0],
             "key": log_key[-1],
+            "method": "PUT",
         }
         if partial:
             params["partial"] = True
-        resp = self._instance.requests_managed_retries_session.post(
+        resp = self._instance.requests_managed_retries_session.get(
             self._instance.dagster_cloud_gen_logs_url_url,
             params=params,
             headers=self._instance.dagster_cloud_api_headers(DagsterCloudInstanceScope.DEPLOYMENT),
@@ -114,12 +120,11 @@ class CloudComputeLogManager(
         if resp_data.get("skip_upload"):
             return
 
-        with open(path, "rb") as f:
-            self._upload_session.post(
-                resp_data["url"],
-                data=resp_data["fields"],
-                files={"file": f},
-            )
+        self._upload_session.put(
+            resp_data["url"],
+            data=data,
+            timeout=self._instance.dagster_cloud_api_timeout,
+        )
 
     def download_from_cloud_storage(
         self, log_key: Sequence[str], io_type: ComputeIOType, partial=False

dagster_cloud/storage/defs_state/__init__.py

@@ -0,0 +1,3 @@
+from dagster_cloud.storage.defs_state.storage import (
+    GraphQLDefsStateStorage as GraphQLDefsStateStorage,
+)

dagster_cloud/storage/defs_state/queries.py

@@ -0,0 +1,15 @@
+GET_LATEST_DEFS_STATE_INFO_QUERY = """
+    query getLatestDefsStateInfo {
+        latestDefsStateInfo
+    }
+"""
+
+SET_LATEST_VERSION_MUTATION = """
+    mutation setLatestDefsStateVersion($key: String!, $version: String!) {
+        defsState {
+            setLatestDefsStateVersion(key: $key, version: $version) {
+                ok
+            }
+        }
+    }
+"""

dagster_cloud/storage/defs_state/storage.py

@@ -0,0 +1,113 @@
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Optional
+
+import dagster._check as check
+from dagster._core.storage.defs_state.base import DefsStateStorage
+from dagster._serdes import ConfigurableClass, ConfigurableClassData
+from dagster_cloud_cli.core.artifacts import download_artifact, upload_artifact
+from dagster_cloud_cli.core.headers.auth import DagsterCloudInstanceScope
+from dagster_shared.serdes import deserialize_value
+from dagster_shared.serdes.objects.models.defs_state_info import DefsStateInfo
+from typing_extensions import Self
+
+if TYPE_CHECKING:
+    from dagster_cloud.instance import DagsterCloudAgentInstance  # noqa: F401
+
+GET_LATEST_DEFS_STATE_INFO_QUERY = """
+    query getLatestDefsStateInfo {
+        latestDefsStateInfo
+    }
+"""
+
+SET_LATEST_VERSION_MUTATION = """
+    mutation setLatestDefsStateVersion($key: String!, $version: String!) {
+        defsState {
+            setLatestDefsStateVersion(key: $key, version: $version) {
+                ok
+            }
+        }
+    }
+"""
+
+
+class GraphQLDefsStateStorage(DefsStateStorage["DagsterCloudAgentInstance"], ConfigurableClass):
+    def __init__(
+        self, inst_data: Optional[ConfigurableClassData] = None, override_graphql_client=None
+    ):
+        """Initialize this class directly only for test (using `override_graphql_client`).
+        Use the ConfigurableClass machinery to init from instance yaml.
+        """
+        self._inst_data = check.opt_inst_param(inst_data, "inst_data", ConfigurableClassData)
+        self._override_graphql_client = override_graphql_client
+
+    @property
+    def inst_data(self):
+        return self._inst_data
+
+    @classmethod
+    def config_type(cls):
+        return {}
+
+    @classmethod
+    def from_config_value(cls, inst_data: ConfigurableClassData, config_value: Any) -> Self:
+        return cls(inst_data=inst_data)
+
+    @property
+    def url(self) -> str:
+        return self._instance.dagster_cloud_url
+
+    @property
+    def api_token(self) -> str:
+        return check.not_none(self._instance.dagster_cloud_agent_token)
+
+    @property
+    def deployment(self) -> str:
+        return check.not_none(self._instance.deployment_name)
+
+    @property
+    def graphql_client(self):
+        return (
+            self._override_graphql_client
+            if self._override_graphql_client
+            else self._instance.graphql_client
+        )
+
+    def _execute_query(self, query, variables=None, idempotent_mutation=False):
+        return self.graphql_client.execute(
+            query, variable_values=variables, idempotent_mutation=idempotent_mutation
+        )
+
+    def _get_artifact_key(self, key: str, version: str) -> str:
+        return f"__state__/{self._sanitize_key(key)}/{version}"
+
+    def download_state_to_path(self, key: str, version: str, path: Path) -> None:
+        download_artifact(
+            url=self.url,
+            scope=DagsterCloudInstanceScope.DEPLOYMENT,
+            api_token=self.api_token,
+            key=self._get_artifact_key(key, version),
+            path=path,
+            deployment=self.deployment,
+        )
+
+    def upload_state_from_path(self, key: str, version: str, path: Path) -> None:
+        upload_artifact(
+            url=self.url,
+            scope=DagsterCloudInstanceScope.DEPLOYMENT,
+            api_token=self.api_token,
+            key=self._get_artifact_key(key, version),
+            path=path,
+            deployment=self.deployment,
+        )
+        self.set_latest_version(key, version)
+
+    def get_latest_defs_state_info(self) -> Optional[DefsStateInfo]:
+        res = self._execute_query(GET_LATEST_DEFS_STATE_INFO_QUERY)
+        result = res["data"]["latestDefsStateInfo"]
+        if result is not None:
+            return deserialize_value(result, DefsStateInfo)
+        else:
+            return None
+
+    def set_latest_version(self, key: str, version: str) -> None:
+        self._execute_query(SET_LATEST_VERSION_MUTATION, variables={"key": key, "version": version})

dagster_cloud/storage/event_logs/__init__.py

@@ -1 +1,3 @@
-from .storage import GraphQLEventLogStorage as GraphQLEventLogStorage
+from dagster_cloud.storage.event_logs.storage import (
+    GraphQLEventLogStorage as GraphQLEventLogStorage,
+)