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

dagster_cloud/__init__.py

@@ -1,7 +1,7 @@
 from dagster._core.libraries import DagsterLibraryRegistry
 
-from .instance import DagsterCloudAgentInstance as DagsterCloudAgentInstance
-from .storage.compute_logs import CloudComputeLogManager as CloudComputeLogManager
-from .version import __version__
+from dagster_cloud.instance import DagsterCloudAgentInstance as DagsterCloudAgentInstance
+from dagster_cloud.storage.compute_logs import CloudComputeLogManager as CloudComputeLogManager
+from dagster_cloud.version import __version__
 
 DagsterLibraryRegistry.register("dagster-cloud", __version__)

dagster_cloud/agent/dagster_cloud_agent.py

@@ -12,8 +12,7 @@ from typing import TYPE_CHECKING, Any, Optional, Union, cast
 import dagster._check as check
 from dagster import DagsterInstance
 from dagster._core.launcher.base import LaunchRunContext
-from dagster._core.remote_representation import CodeLocationOrigin
-from dagster._core.remote_representation.origin import RegisteredCodeLocationOrigin
+from dagster._core.remote_origin import CodeLocationOrigin, RegisteredCodeLocationOrigin
 from dagster._core.utils import FuturesAwareThreadPoolExecutor
 from dagster._grpc.client import DagsterGrpcClient
 from dagster._grpc.types import CancelExecutionRequest
@@ -21,13 +20,27 @@ from dagster._serdes import deserialize_value, serialize_value
 from dagster._time import get_current_datetime, get_current_timestamp
 from dagster._utils.cached_method import cached_method
 from dagster._utils.container import retrieve_containerized_utilization_metrics
-from dagster._utils.error import SerializableErrorInfo, serializable_error_info_from_exc_info
+from dagster._utils.error import (
+    SerializableErrorInfo,
+    serializable_error_info_from_exc_info,
+    truncate_serialized_error,
+)
 from dagster._utils.interrupts import raise_interrupts_as
 from dagster._utils.merger import merge_dicts
 from dagster._utils.typed_dict import init_optional_typeddict
 from dagster_cloud_cli.core.errors import DagsterCloudHTTPError, raise_http_error
 from dagster_cloud_cli.core.workspace import CodeLocationDeployData
 
+from dagster_cloud.agent.instrumentation.constants import DAGSTER_CLOUD_AGENT_METRIC_PREFIX
+from dagster_cloud.agent.instrumentation.run_launch import extract_run_attributes
+from dagster_cloud.agent.instrumentation.schedule import inspect_schedule_result
+from dagster_cloud.agent.instrumentation.sensor import inspect_sensor_result
+from dagster_cloud.agent.queries import (
+    ADD_AGENT_HEARTBEATS_MUTATION,
+    DEPLOYMENTS_QUERY,
+    GET_USER_CLOUD_REQUESTS_QUERY,
+    WORKSPACE_ENTRIES_QUERY,
+)
 from dagster_cloud.api.dagster_cloud_api import (
     AgentHeartbeat,
     AgentUtilizationMetrics,
@@ -45,27 +58,15 @@ from dagster_cloud.api.dagster_cloud_api import (
 )
 from dagster_cloud.batching import Batcher
 from dagster_cloud.instance import DagsterCloudAgentInstance
-from dagster_cloud.util.errors import truncate_serialized_error
+from dagster_cloud.opentelemetry.observers.execution_observer import observe_execution
+from dagster_cloud.util import SERVER_HANDLE_TAG, compressed_namedtuple_upload_file, is_isolated_run
+from dagster_cloud.version import __version__
 from dagster_cloud.workspace.user_code_launcher import (
     DagsterCloudUserCodeLauncher,
     UserCodeLauncherEntry,
 )
 from dagster_cloud.workspace.user_code_launcher.utils import get_instance_ref_for_user_code
 
-from ..opentelemetry.observers.execution_observer import observe_execution
-from ..util import SERVER_HANDLE_TAG, compressed_namedtuple_upload_file, is_isolated_run
-from ..version import __version__
-from .instrumentation.constants import DAGSTER_CLOUD_AGENT_METRIC_PREFIX
-from .instrumentation.run_launch import extract_run_attributes
-from .instrumentation.schedule import inspect_schedule_result
-from .instrumentation.sensor import inspect_sensor_result
-from .queries import (
-    ADD_AGENT_HEARTBEATS_MUTATION,
-    DEPLOYMENTS_QUERY,
-    GET_USER_CLOUD_REQUESTS_QUERY,
-    WORKSPACE_ENTRIES_QUERY,
-)
-
 if TYPE_CHECKING:
     import datetime
 
@@ -162,6 +163,8 @@ class DagsterCloudAgent:
 
         self._last_liveness_check_time = None
 
+        self._warned_about_long_in_progress_reconcile = False
+
     def __enter__(self):
         return self
 
@@ -305,6 +308,19 @@ class DagsterCloudAgent:
 
         self._check_initial_deployment_names()
 
+        serving = []
+        queues = list(filter(None, self._instance.agent_queues_config.queues))
+        if queues:
+            serving.append(f"queues{queues}")
+        if self._instance.deployment_names:
+            serving.append(f"deployments{self._instance.deployment_names}")
+        if self._instance.include_all_serverless_deployments:
+            serving.append("all serverless deployments")
+        if self._instance.includes_branch_deployments:
+            serving.append("branch deployments")
+
+        self._logger.info(f"Agent is serving: {', '.join(serving)}")
+
         self._check_update_workspace(
             user_code_launcher,
             upload_all=user_code_launcher.upload_snapshots_on_startup,
@@ -352,6 +368,8 @@ class DagsterCloudAgent:
                 except Exception:
                     self._logger.exception("Failed to add heartbeat")
 
+            self._check_for_long_running_reconcile(user_code_launcher)
+
             # Check for any received interrupts
             with raise_interrupts_as(KeyboardInterrupt):
                 pass
@@ -394,6 +412,26 @@ class DagsterCloudAgent:
             self._logger.error(f"Failed to write liveness sentinel and disabling it: {e}")
             self._last_liveness_check_time = False
 
+    def _check_for_long_running_reconcile(self, user_code_launcher):
+        """Detect from the main thread if the background reconcile thread is running behind or has gotten stuck."""
+        in_progress_reconcile_start_time = user_code_launcher.in_progress_reconcile_start_time
+
+        reconcile_start_time_warning = int(
+            os.getenv("DAGSTER_CLOUD_AGENT_RECONCILE_START_TIME_WARNING", "3600")
+        )
+
+        if (
+            in_progress_reconcile_start_time is not None
+            and (time.time() - in_progress_reconcile_start_time) >= reconcile_start_time_warning
+        ):
+            if not self._warned_about_long_in_progress_reconcile:
+                self._logger.warning(
+                    f"Agent has been redeploying code servers for more than {reconcile_start_time_warning} seconds. This may indicate the background thread that performs the redeploys is stuck."
+                )
+                self._warned_about_long_in_progress_reconcile = True
+        else:
+            self._warned_about_long_in_progress_reconcile = False
+
     def _check_update_workspace(self, user_code_launcher, upload_all):
         curr_time = get_current_datetime()
 
@@ -1044,10 +1082,16 @@ class DagsterCloudAgent:
                         client.cancel_execution(CancelExecutionRequest(run_id=run.run_id))
 
             return DagsterCloudApiSuccess()
-
+        elif api_name in (
+            DagsterCloudApi.CHECK_STEP_HEALTH,
+            DagsterCloudApi.TERMINATE_STEP,
+            DagsterCloudApi.LAUNCH_STEP,
+            DagsterCloudApi.CHECK_RUN_HEALTH,
+            DagsterCloudApi.LOAD_REPOSITORIES,
+        ):
+            check.failed(f"Unexpected deprecated request type {api_name}")
         else:
             check.assert_never(api_name)
-            raise Exception(f"Unexpected dagster cloud api call {api_name}")
 
     def _process_api_request(
         self,

dagster_cloud/agent/instrumentation/run_launch.py

@@ -19,5 +19,5 @@ def extract_run_attributes(deployment_name, run: DagsterRun) -> dict[str, str]:
             attributes["location"] = run.remote_job_origin.location_name or "unknown"
     except Exception as e:  # pylint: disable=broad-except
         logging.warning("Failed to extract additional run attributes", exc_info=e)
-    finally:
-        return attributes
+
+    return attributes

dagster_cloud/agent/instrumentation/schedule.py

@@ -1,6 +1,6 @@
 from dagster import deserialize_value
 from dagster._core.definitions.schedule_definition import ScheduleExecutionData
-from dagster._core.remote_representation import ScheduleExecutionErrorSnap
+from dagster._core.remote_representation.external_data import ScheduleExecutionErrorSnap
 
 from dagster_cloud.opentelemetry.observers.execution_observer import ExecutionObserverInstruments
 

dagster_cloud/agent/instrumentation/sensor.py

@@ -1,6 +1,6 @@
 from dagster import deserialize_value
 from dagster._core.definitions.sensor_definition import SensorExecutionData
-from dagster._core.remote_representation import SensorExecutionErrorSnap
+from dagster._core.remote_representation.external_data import SensorExecutionErrorSnap
 
 from dagster_cloud.opentelemetry.observers.execution_observer import ExecutionObserverInstruments
 

dagster_cloud/anomaly_detection/__init__.py

@@ -1,7 +1,7 @@
-from .defs import (
+from dagster_cloud.anomaly_detection.defs import (
     build_anomaly_detection_freshness_checks as build_anomaly_detection_freshness_checks,
 )
-from .types import (
+from dagster_cloud.anomaly_detection.types import (
     AnomalyDetectionModelParams as AnomalyDetectionModelParams,
     BetaFreshnessAnomalyDetectionParams as BetaFreshnessAnomalyDetectionParams,
 )

dagster_cloud/anomaly_detection/defs.py

@@ -7,24 +7,27 @@ from dagster import (
     MetadataValue,
     _check as check,
 )
-from dagster._core.definitions.asset_check_factories.freshness_checks.last_update import (
+from dagster._core.definitions.asset_checks.asset_check_factories.freshness_checks.last_update import (
     construct_description,
 )
-from dagster._core.definitions.asset_check_factories.utils import (
+from dagster._core.definitions.asset_checks.asset_check_factories.utils import (
     FRESH_UNTIL_METADATA_KEY,
     LAST_UPDATED_TIMESTAMP_METADATA_KEY,
     LOWER_BOUND_TIMESTAMP_METADATA_KEY,
     assets_to_keys,
 )
-from dagster._core.definitions.asset_check_result import AssetCheckResult
-from dagster._core.definitions.asset_check_spec import (
+from dagster._core.definitions.asset_checks.asset_check_result import AssetCheckResult
+from dagster._core.definitions.asset_checks.asset_check_spec import (
     AssetCheckKey,
     AssetCheckSeverity,
     AssetCheckSpec,
 )
-from dagster._core.definitions.asset_checks import AssetChecksDefinition
+from dagster._core.definitions.asset_checks.asset_checks_definition import AssetChecksDefinition
 from dagster._core.definitions.asset_key import AssetKey
-from dagster._core.definitions.assets import AssetsDefinition, unique_id_from_asset_and_check_keys
+from dagster._core.definitions.assets.definition.assets_definition import (
+    AssetsDefinition,
+    unique_id_from_asset_and_check_keys,
+)
 from dagster._core.definitions.decorators.asset_check_decorator import multi_asset_check
 from dagster._core.definitions.events import CoercibleToAssetKey
 from dagster._core.definitions.source_asset import SourceAsset
@@ -36,8 +39,8 @@ from dagster_cloud_cli.core.graphql_client import (
     create_cloud_webserver_client,
 )
 
-from .mutation import ANOMALY_DETECTION_INFERENCE_MUTATION
-from .types import (
+from dagster_cloud.anomaly_detection.mutation import ANOMALY_DETECTION_INFERENCE_MUTATION
+from dagster_cloud.anomaly_detection.types import (
     AnomalyDetectionModelParams,
     BetaFreshnessAnomalyDetectionParams,
     FreshnessAnomalyDetectionResult,

dagster_cloud/api/dagster_cloud_api.py

@@ -7,7 +7,8 @@ from typing import Any, Optional, TypedDict, Union
 from dagster._core.code_pointer import CodePointer
 from dagster._core.definitions.selector import JobSelector
 from dagster._core.events.log import EventLogEntry
-from dagster._core.remote_representation import CodeLocationOrigin, RepositorySnap
+from dagster._core.remote_origin import CodeLocationOrigin
+from dagster._core.remote_representation.external_data import RepositorySnap
 from dagster._core.storage.dagster_run import DagsterRun
 from dagster._core.utils import RequestUtilizationMetrics
 from dagster._record import IHaveNew, copy, record, record_custom
@@ -84,7 +85,6 @@ class DagsterCloudUploadWorkspaceResponse:
 @whitelist_for_serdes
 class DagsterCloudApi(Enum):
     CHECK_FOR_WORKSPACE_UPDATES = "CHECK_FOR_WORKSPACE_UPDATES"
-    LOAD_REPOSITORIES = "LOAD_REPOSITORIES"
     GET_EXTERNAL_EXECUTION_PLAN = "GET_EXTERNAL_EXECUTION_PLAN"
     GET_SUBSET_EXTERNAL_PIPELINE_RESULT = "GET_SUBSET_EXTERNAL_PIPELINE_RESULT"
     GET_EXTERNAL_PARTITION_CONFIG = "GET_EXTERNAL_PARTITION_CONFIG"
@@ -98,13 +98,15 @@ class DagsterCloudApi(Enum):
     GET_EXTERNAL_NOTEBOOK_DATA = "GET_EXTERNAL_NOTEBOOK_DATA"
 
     LAUNCH_RUN = "LAUNCH_RUN"
-    CHECK_RUN_HEALTH = "CHECK_RUN_HEALTH"  # deprecated, agents now surface this in heartbeats
     TERMINATE_RUN = "TERMINATE_RUN"
+
+    PING_LOCATION = "PING_LOCATION"  # Signal that a location is in use and should keep servers up
+
+    CHECK_RUN_HEALTH = "CHECK_RUN_HEALTH"  # deprecated, agents now surface this in heartbeats
     LAUNCH_STEP = "LAUNCH_STEP"  # deprecated with cloud executor
     CHECK_STEP_HEALTH = "CHECK_STEP_HEALTH"  # deprecated with cloud executor
     TERMINATE_STEP = "TERMINATE_STEP"  # deprecated with cloud executor
-
-    PING_LOCATION = "PING_LOCATION"  # Signal that a location is in use and should keep servers up
+    LOAD_REPOSITORIES = "LOAD_REPOSITORIES"  # deprecated
 
     def __structlog__(self):
         return self.name

dagster_cloud/auth/constants.py

@@ -1,3 +1,4 @@
+import uuid
 from typing import Optional
 
 from dagster._core.errors import DagsterInvariantViolationError
@@ -19,12 +20,27 @@ def get_organization_public_id_from_api_token(api_token: str) -> Optional[str]:
     return split_token[2]
 
 
-def get_organization_name_from_agent_token(agent_token: str) -> Optional[str]:
+def decode_region_from_uuid(regional_token: str) -> Optional[str]:
+    try:
+        regional_uuid = uuid.UUID(regional_token)
+    except ValueError:
+        # if it's not an actual uuid, we can't decode region
+        return None
+
+    # custom uuids contain region subdomains in the first 2 bytes
+    if regional_uuid.version != 8 or regional_uuid.variant != uuid.RFC_4122:
+        return None
+
+    uuid_bytes = regional_uuid.bytes
+    return uuid_bytes[:2].decode("ascii")
+
+
+def decode_agent_token(agent_token: str) -> tuple[Optional[str], Optional[str]]:
     split_token = agent_token.split(":")
 
     # Legacy agent token format - organization must be specified in dagster.yaml
     if len(split_token) == 1:
-        return None
+        return None, None
 
     token_type, *token = split_token
 
@@ -35,6 +51,6 @@ def get_organization_name_from_agent_token(agent_token: str) -> Optional[str]:
             "Generate a new agent token in Dagster Cloud."
         )
 
-    organization, _identifier = token
-
-    return organization
+    # token format: agent:<org>:<uuid>
+    organization, uuid_str = token
+    return organization, decode_region_from_uuid(uuid_str)

dagster_cloud/batching/__init__.py

@@ -1 +1 @@
-from .batcher import Batcher as Batcher
+from dagster_cloud.batching.batcher import Batcher as Batcher

dagster_cloud/dagster_insights/__init__.py

@@ -1,15 +1,19 @@
 import sys
 from typing import Any
 
-from .snowflake.dbt_wrapper import dbt_with_snowflake_insights as dbt_with_snowflake_insights
-from .snowflake.definitions import (
+from dagster_cloud.dagster_insights.snowflake.dbt_wrapper import (
+    dbt_with_snowflake_insights as dbt_with_snowflake_insights,
+)
+from dagster_cloud.dagster_insights.snowflake.definitions import (
     create_snowflake_insights_asset_and_schedule as create_snowflake_insights_asset_and_schedule,
 )
-from .snowflake.snowflake_utils import meter_snowflake_query as meter_snowflake_query
+from dagster_cloud.dagster_insights.snowflake.snowflake_utils import (
+    meter_snowflake_query as meter_snowflake_query,
+)
 
 dagster_snowflake_req_imports = {"InsightsSnowflakeResource"}
 try:
-    from .snowflake.insights_snowflake_resource import (
+    from dagster_cloud.dagster_insights.snowflake.insights_snowflake_resource import (
         InsightsSnowflakeResource as InsightsSnowflakeResource,
     )
 except ImportError:
@@ -17,8 +21,10 @@ except ImportError:
 
 dagster_bigquery_req_imports = {"InsightsBigQueryResource", "dbt_with_bigquery_insights"}
 try:
-    from .bigquery.dbt_wrapper import dbt_with_bigquery_insights as dbt_with_bigquery_insights
-    from .bigquery.insights_bigquery_resource import (
+    from dagster_cloud.dagster_insights.bigquery.dbt_wrapper import (
+        dbt_with_bigquery_insights as dbt_with_bigquery_insights,
+    )
+    from dagster_cloud.dagster_insights.bigquery.insights_bigquery_resource import (
         InsightsBigQueryResource as InsightsBigQueryResource,
     )
 except ImportError:

dagster_cloud/dagster_insights/bigquery/dbt_wrapper.py

@@ -18,8 +18,14 @@ from dagster_dbt import DbtCliInvocation
 from dagster_dbt.version import __version__ as dagster_dbt_version
 from packaging import version
 
-from ..insights_utils import extract_asset_info_from_event, handle_raise_on_error
-from .bigquery_utils import build_bigquery_cost_metadata, marker_asset_key_for_job
+from dagster_cloud.dagster_insights.bigquery.bigquery_utils import (
+    build_bigquery_cost_metadata,
+    marker_asset_key_for_job,
+)
+from dagster_cloud.dagster_insights.insights_utils import (
+    extract_asset_info_from_event,
+    handle_raise_on_error,
+)
 
 if TYPE_CHECKING:
     from dbt.adapters.base.impl import BaseAdapter

dagster_cloud/dagster_insights/bigquery/insights_bigquery_resource.py

@@ -7,10 +7,12 @@ from dagster_gcp import BigQueryResource
 from dagster_gcp.bigquery.utils import setup_gcp_creds
 from google.cloud import bigquery
 
+from dagster_cloud.dagster_insights.bigquery.bigquery_utils import (
+    build_bigquery_cost_metadata,
+    marker_asset_key_for_job,
+)
 from dagster_cloud.dagster_insights.insights_utils import get_current_context_and_asset_key
 
-from .bigquery_utils import build_bigquery_cost_metadata, marker_asset_key_for_job
-
 OUTPUT_NON_ASSET_SIGIL = "__bigquery_query_metadata_"
 
 

dagster_cloud/dagster_insights/insights_utils.py

@@ -33,7 +33,7 @@ def get_current_context_and_asset_key() -> tuple[
 def get_asset_key_for_output(
     context: Union[OpExecutionContext, AssetExecutionContext], output_name: str
 ) -> Optional[AssetKey]:
-    asset_key = context.job_def.asset_layer.asset_key_for_output(
+    asset_key = context.job_def.asset_layer.get_asset_key_for_node_output(
         node_handle=context.op_handle, output_name=output_name
     )
     if asset_key is None:

dagster_cloud/dagster_insights/metrics_utils.py

@@ -28,7 +28,7 @@ def get_url_and_token_from_instance(instance: DagsterInstance) -> tuple[str, str
     if not isinstance(instance, DagsterCloudAgentInstance):
         raise RuntimeError("This asset only functions in a running Dagster Cloud instance")
 
-    return f"{instance.dagit_url}graphql", instance.dagster_cloud_agent_token  # pyright: ignore[reportReturnType]
+    return f"{instance.dagit_url}graphql", instance.dagster_cloud_agent_token
 
 
 def get_insights_upload_request_params(

dagster_cloud/dagster_insights/snowflake/dagster_snowflake_insights.py

@@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Optional
 
 from dagster import AssetKey, AssetsDefinition, ScheduleDefinition
 
-from .snowflake_utils import OPAQUE_ID_SQL_SIGIL
+from dagster_cloud.dagster_insights.snowflake.snowflake_utils import OPAQUE_ID_SQL_SIGIL
 
 if TYPE_CHECKING:
     from dagster_snowflake import SnowflakeConnection
@@ -78,10 +78,6 @@ HAVING ARRAY_SIZE(opaque_ids) > 0
 
     costs: list[tuple[str, float, str]] = []
 
-    print(
-        f"{len(results) if results else 0} annotated queries returned from snowflake query_history"
-    )
-
     if not results:
         return []
 
@@ -94,8 +90,4 @@ HAVING ARRAY_SIZE(opaque_ids) > 0
         for opaque_id in opaque_ids:
             costs.append((opaque_id, float(cost), query_id))
 
-    print(
-        f"Reported costs for {len(costs)} of {total} asset materializations found in the"
-        " query_history."
-    )
     return costs

dagster_cloud/dagster_insights/snowflake/dbt_wrapper.py

@@ -13,8 +13,15 @@ from dagster import (
     Output,
 )
 
-from ..insights_utils import extract_asset_info_from_event, handle_raise_on_error
-from .snowflake_utils import OPAQUE_ID_SQL_SIGIL, build_opaque_id_metadata, marker_asset_key_for_job
+from dagster_cloud.dagster_insights.insights_utils import (
+    extract_asset_info_from_event,
+    handle_raise_on_error,
+)
+from dagster_cloud.dagster_insights.snowflake.snowflake_utils import (
+    OPAQUE_ID_SQL_SIGIL,
+    build_opaque_id_metadata,
+    marker_asset_key_for_job,
+)
 
 if TYPE_CHECKING:
     from dagster_dbt import DbtCliInvocation

dagster_cloud/dagster_insights/snowflake/definitions.py

@@ -2,7 +2,6 @@ import warnings
 from collections.abc import Sequence
 from dataclasses import dataclass
 from datetime import date, datetime, timedelta, timezone
-from pprint import pprint
 from typing import TYPE_CHECKING, Optional, Union
 
 from dagster import (
@@ -27,8 +26,10 @@ from dagster._core.storage.tags import (
     ASSET_PARTITION_RANGE_START_TAG,
 )
 
-from ..metrics_utils import put_cost_information
-from .dagster_snowflake_insights import get_cost_data_for_hour
+from dagster_cloud.dagster_insights.metrics_utils import put_cost_information
+from dagster_cloud.dagster_insights.snowflake.dagster_snowflake_insights import (
+    get_cost_data_for_hour,
+)
 
 if TYPE_CHECKING:
     from dagster_snowflake import SnowflakeConnection
@@ -154,7 +155,7 @@ def create_snowflake_insights_asset_and_schedule(
         )
 
         if dry_run:
-            pprint(costs)
+            pass
         else:
             context.log.info(
                 f"Submitting cost information for {len(costs)} queries to Dagster Insights"

dagster_cloud/dagster_insights/snowflake/insights_snowflake_resource.py

@@ -20,8 +20,7 @@ from dagster_snowflake import SnowflakeConnection, SnowflakeResource
 from snowflake.connector.cursor import SnowflakeCursor
 
 from dagster_cloud.dagster_insights.insights_utils import get_current_context_and_asset_key
-
-from .snowflake_utils import meter_snowflake_query
+from dagster_cloud.dagster_insights.snowflake.snowflake_utils import meter_snowflake_query
 
 
 def get_current_context_and_asset_key_or_warn() -> tuple[

dagster_cloud/definitions/job_selection.py

@@ -0,0 +1,36 @@
+from abc import ABC, abstractmethod
+from collections.abc import Sequence
+from typing import AbstractSet  # noqa: UP035
+
+import dagster._check as check
+from dagster._record import record
+from dagster_shared.serdes import whitelist_for_serdes
+
+
+@whitelist_for_serdes
+@record
+class JobSelection(ABC):
+    code_location_name: str
+    repository_name: str
+
+    @abstractmethod
+    def resolve_job_names(self) -> AbstractSet[str]:
+        raise NotImplementedError()
+
+    @staticmethod
+    def names(
+        names: Sequence[str], code_location_name: str, repository_name: str
+    ) -> "NamesJobSelection":
+        check.invariant(len(names) == 1, "Only one job name is supported at this time")
+        return NamesJobSelection(
+            job_names=names, code_location_name=code_location_name, repository_name=repository_name
+        )
+
+
+@whitelist_for_serdes
+@record
+class NamesJobSelection(JobSelection):
+    job_names: Sequence[str]
+
+    def resolve_job_names(self) -> AbstractSet[str]:
+        return set(self.job_names)

dagster_cloud/execution/utils/process.py

@@ -4,7 +4,7 @@ from collections.abc import Sequence
 
 from dagster_shared.ipc import interrupt_ipc_subprocess_pid, open_ipc_subprocess
 
-from . import TaskStatus
+from dagster_cloud.execution.utils import TaskStatus
 
 
 def launch_process(args: Sequence[str]) -> int: