dagster-cloud 1.8.2__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/__init__.py

@@ -1,4 +1,4 @@
-from typing import List, NamedTuple, Optional
+from typing import NamedTuple, Optional
 
 from dagster._serdes import whitelist_for_serdes
 from dagster_cloud_cli.core.agent_queue import AgentQueue
@@ -7,14 +7,14 @@ from dagster_cloud_cli.core.agent_queue import AgentQueue
 @whitelist_for_serdes
 class AgentQueuesConfig(NamedTuple):
     include_default_queue: bool = True
-    additional_queues: Optional[List[AgentQueue]] = None
+    additional_queues: Optional[list[AgentQueue]] = None
 
     @staticmethod
-    def default_queues() -> List[AgentQueue]:
+    def default_queues() -> list[AgentQueue]:
         return [None]
 
     @property
-    def queues(self) -> List[AgentQueue]:
+    def queues(self) -> list[AgentQueue]:
         if self.additional_queues:
             queues = self.additional_queues.copy()
         else:

dagster_cloud/agent/cli/__init__.py

@@ -1,17 +1,20 @@
 import json
 import logging
+import logging.config
 import os
+from collections.abc import Mapping
 from pathlib import Path
 from tempfile import TemporaryDirectory
-from typing import Optional
+from typing import Optional, cast
 
+import dagster._check as check
 import yaml
 from dagster._core.errors import DagsterHomeNotSetError
 from dagster._utils.interrupts import capture_interrupts
 from dagster._utils.log import default_date_format_string, default_format_string
 from dagster._utils.merger import deep_merge_dicts
-from dagster._utils.yaml_utils import load_yaml_from_globs
 from dagster_cloud_cli import ui
+from dagster_shared.yaml_utils import load_yaml_from_globs
 from typer import Argument, Option, Typer
 
 from dagster_cloud.agent.dagster_cloud_agent import DagsterCloudAgent
@@ -35,9 +38,11 @@ def agent_home_exception():
     return ui.error(f"No Dagster config found.\n\n{dagster_home_loc}")
 
 
-def run_local_agent(agent_logging_config_path):
+def run_local_agent(agent_logging_config: Optional[Mapping[str, object]]) -> None:
     try:
-        with DagsterCloudAgentInstance.get() as instance:
+        with DagsterCloudAgentInstance.get() as inst:
+            instance = check.inst(inst, DagsterCloudAgentInstance)
+
             logging.basicConfig(
                 level=logging.INFO,
                 format=default_format_string(),
@@ -45,25 +50,26 @@ def run_local_agent(agent_logging_config_path):
                 handlers=[logging.StreamHandler()],
             )
 
-            if agent_logging_config_path:
-                agent_logging_config = {
-                    "version": 1,
-                    "disable_existing_loggers": False,
-                    **load_yaml_from_globs(str(agent_logging_config_path)),
-                }
-                logging.config.dictConfig(agent_logging_config)
+            if agent_logging_config:
+                logging.config.dictConfig(
+                    {
+                        "version": 1,
+                        "disable_existing_loggers": False,
+                        **agent_logging_config,
+                    }
+                )
 
             user_code_launcher = instance.user_code_launcher
             user_code_launcher.start()
 
-            with DagsterCloudAgent() as agent:
-                agent.run_loop(instance, user_code_launcher, agent_uuid=instance.instance_uuid)
+            with DagsterCloudAgent(instance) as agent:
+                agent.run_loop(user_code_launcher, agent_uuid=instance.instance_uuid)
     except DagsterHomeNotSetError:
         raise agent_home_exception()
 
 
 def run_local_agent_in_environment(
-    dagster_home: Optional[Path], agent_logging_config_path: Optional[Path]
+    dagster_home: Optional[Path], agent_logging_config: Optional[Mapping[str, object]]
 ):
     with capture_interrupts():
         old_env = None
@@ -71,7 +77,7 @@ def run_local_agent_in_environment(
             old_env = dict(os.environ)
             if dagster_home:
                 os.environ["DAGSTER_HOME"] = str(dagster_home.resolve())
-            run_local_agent(agent_logging_config_path)
+            run_local_agent(agent_logging_config)
         finally:
             os.environ.clear()
             if old_env is not None:
@@ -86,6 +92,7 @@ def run_local_agent_in_temp_environment(
     user_code_launcher_module: Optional[str],
     user_code_launcher_class: Optional[str],
     user_code_launcher_config: Optional[str],
+    agent_logging_config: Optional[Mapping[str, object]],
 ):
     config = {
         "instance_class": {
@@ -121,7 +128,7 @@ def run_local_agent_in_temp_environment(
     with TemporaryDirectory() as d:
         with open(os.path.join(d, "dagster.yaml"), "w", encoding="utf8") as f:
             f.write(yaml.dump(config))
-        run_local_agent_in_environment(Path(d), None)
+        run_local_agent_in_environment(Path(d), agent_logging_config)
 
 
 @app.command(
@@ -170,7 +177,23 @@ def run(
         ),
         exists=True,
     ),
+    agent_logging_config_string: Optional[str] = Option(
+        None,
+        "--agent-logging-config-string",
+        help=(
+            "inlined json with logging config for the agent process that can be passed into"
+            " logging.dictConfig. Cannot be provided if --agent-logging-config-path is specified."
+        ),
+    ),
 ):
+    if agent_logging_config_string and agent_logging_config_path:
+        raise ui.error(
+            "Only --agent-logging-config-path or --agent-logging-config-string can be specified, not both"
+        )
+    else:
+        agent_logging_config = _get_agent_logging_config(
+            agent_logging_config_path, agent_logging_config_string
+        )
     if (
         agent_token
         or deployment
@@ -203,6 +226,22 @@ def run(
             user_code_launcher_module,
             user_code_launcher_class,
             user_code_launcher_config,
+            agent_logging_config,
         )
     else:
-        run_local_agent_in_environment(dagster_home, agent_logging_config_path)
+        run_local_agent_in_environment(dagster_home, agent_logging_config)
+
+
+def _get_agent_logging_config(
+    agent_logging_config_path: Optional[Path],
+    agent_logging_config_string: Optional[str],
+) -> Optional[Mapping[str, object]]:
+    agent_logging_config: Optional[Mapping[str, object]] = None
+    if agent_logging_config_path:
+        agent_logging_config = cast(
+            "Mapping[str, object]",
+            load_yaml_from_globs(str(agent_logging_config_path)),
+        )
+    elif agent_logging_config_string:
+        agent_logging_config = json.loads(agent_logging_config_string)
+    return agent_logging_config