dagster-azure 0.27.15__tar.gz → 0.28.0__tar.gz

{dagster_azure-0.27.15/dagster_azure.egg-info → dagster_azure-0.28.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dagster-azure
-Version: 0.27.15
+Version: 0.28.0
 Summary: Package for Azure-specific Dagster framework op and resource components.
 Home-page: https://github.com/dagster-io/dagster/tree/master/python_modules/libraries/dagster-azure
 Author: Dagster Labs
@@ -17,9 +17,10 @@ Requires-Python: >=3.9,<3.14
 License-File: LICENSE
 Requires-Dist: azure-core<2.0.0,>=1.7.0
 Requires-Dist: azure-identity<2.0.0,>=1.7.0
+Requires-Dist: azure-ai-ml<2.0.0,>=1.28.0
 Requires-Dist: azure-storage-blob<13.0.0,>=12.5.0
 Requires-Dist: azure-storage-file-datalake<13.0.0,>=12.5
-Requires-Dist: dagster==1.11.15
+Requires-Dist: dagster==1.12.0
 Dynamic: author
 Dynamic: author-email
 Dynamic: classifier

dagster_azure-0.28.0/dagster_azure/pipes/__init__.py

@@ -0,0 +1,9 @@
+from dagster_azure.pipes.clients import PipesAzureMLClient
+from dagster_azure.pipes.context_injectors import PipesAzureBlobStorageContextInjector
+from dagster_azure.pipes.message_readers import PipesAzureBlobStorageMessageReader
+
+__all__ = [
+    "PipesAzureBlobStorageContextInjector",
+    "PipesAzureBlobStorageMessageReader",
+    "PipesAzureMLClient",
+]

dagster_azure-0.28.0/dagster_azure/pipes/clients/__init__.py

@@ -0,0 +1,5 @@
+from dagster_azure.pipes.clients.azureml import PipesAzureMLClient
+
+__all__ = [
+    "PipesAzureMLClient",
+]

dagster_azure-0.28.0/dagster_azure/pipes/clients/azureml.py

@@ -0,0 +1,140 @@
+import time
+from typing import Optional, Union
+
+import dagster._check as check
+from azure.ai.ml import MLClient
+from azure.ai.ml.entities import Command
+from dagster._core.definitions.resource_annotation import TreatAsResourceParam
+from dagster._core.errors import DagsterExecutionInterruptedError, DagsterPipesExecutionError
+from dagster._core.execution.context.asset_execution_context import AssetExecutionContext
+from dagster._core.execution.context.op_execution_context import OpExecutionContext
+from dagster._core.pipes.client import (
+    PipesClient,
+    PipesClientCompletedInvocation,
+    PipesContextInjector,
+    PipesMessageReader,
+)
+from dagster._core.pipes.utils import open_pipes_session
+from dagster_pipes import PipesExtras
+
+
+class PipesAzureMLClient(PipesClient, TreatAsResourceParam):
+    """Pipes client for Azure ML.
+
+    Args:
+        client (MLClient): An Azure ML `MLClient` object.
+        context_injector (PipesContextInjector): A context injector to use to inject
+            context into the Azure ML job process.
+        message_reader (PipesMessageReader): A message reader to use to read messages
+            from the Azure ML job.
+        poll_interval_seconds (float): How long to sleep between checking the status of the job run.
+            Defaults to 5.
+        forward_termination (bool): Whether to cancel the Azure ML job if the orchestration process
+            is interrupted or canceled. Defaults to True.
+    """
+
+    def __init__(
+        self,
+        client: MLClient,
+        context_injector: PipesContextInjector,
+        message_reader: PipesMessageReader,
+        poll_interval_seconds: float = 5,
+        forward_termination: bool = True,
+    ):
+        self.client = client
+        self.context_injector = check.inst_param(
+            context_injector,
+            "context_injector",
+            PipesContextInjector,
+        )
+        self.message_reader = check.inst_param(
+            message_reader,
+            "message_reader",
+            PipesMessageReader,
+        )
+        self.poll_interval_seconds = check.numeric_param(
+            poll_interval_seconds, "poll_interval_seconds"
+        )
+        self.forward_termination = check.bool_param(forward_termination, "forward_termination")
+
+    def _poll_til_success(
+        self, context: Union[OpExecutionContext, AssetExecutionContext], job_name: str
+    ) -> None:
+        # poll the Azure ML job until it completes successfully, raising otherwise
+
+        last_observed_status = None
+
+        while True:
+            status = self.client.jobs.get(job_name).status
+
+            if status != last_observed_status:
+                context.log.info(
+                    f"[pipes] Azure ML job {job_name} observed state transition to {status}"
+                )
+            last_observed_status = status
+
+            if status == "Completed":
+                return
+            elif status in {"Canceled", "Failed"}:
+                raise DagsterPipesExecutionError(f"Error running Azure ML job: {status}")
+
+            time.sleep(self.poll_interval_seconds)
+
+    def _poll_til_terminating(self, job_name: str) -> None:
+        while True:
+            status = self.client.jobs.get(job_name).status
+            if status in {"Completed", "Canceled", "Failed"}:
+                return
+
+            time.sleep(self.poll_interval_seconds)
+
+    @classmethod
+    def _is_dagster_maintained(cls) -> bool:
+        return True
+
+    def run(  # pyright: ignore[reportIncompatibleMethodOverride]
+        self,
+        *,
+        context: Union[OpExecutionContext, AssetExecutionContext],
+        extras: Optional[PipesExtras] = None,
+        command: Command,
+    ) -> PipesClientCompletedInvocation:
+        """Synchronously execute an Azure ML job with the pipes protocol.
+
+        Args:
+            command (azure.ai.ml.entities.Command): Specification of the Azure ML
+                command job to run. Pipes bootstrap parameters will be passed via environment
+                variables which will be merged with any existing environment variables in the command.
+                The Azure ML job code should use :py:func:`dagster_pipes.open_dagster_pipes` to
+                initialize the Pipes connection.
+            context (Union[OpExecutionContext, AssetExecutionContext]): The context from the executing op or asset.
+            extras (Optional[PipesExtras]): An optional dict of extra parameters to pass to the
+                Azure ML job.
+
+        Returns:
+            PipesClientCompletedInvocation: Wrapper containing results reported by the external
+                process.
+        """
+        with open_pipes_session(
+            context=context,
+            extras=extras,
+            context_injector=self.context_injector,
+            message_reader=self.message_reader,
+        ) as pipes_session:
+            command.environment_variables = {
+                **command.environment_variables,
+                **pipes_session.get_bootstrap_env_vars(),
+            }
+            job = self.client.create_or_update(command)
+
+            try:
+                self._poll_til_success(context, job.name)  # pyright: ignore[reportArgumentType]
+            except DagsterExecutionInterruptedError:
+                if self.forward_termination:
+                    context.log.info("[pipes] execution interrupted, canceling Azure ML job.")
+                    self.client.jobs.begin_cancel(job.name)  # pyright: ignore[reportArgumentType]
+                    self._poll_til_terminating(job.name)  # pyright: ignore[reportArgumentType]
+
+        return PipesClientCompletedInvocation(
+            pipes_session, metadata={"AzureML Job Name": job.name}
+        )

dagster_azure-0.28.0/dagster_azure/pipes/context_injectors.py

@@ -0,0 +1,47 @@
+import json
+import os
+import random
+import string
+from collections.abc import Iterator
+from contextlib import contextmanager
+
+import dagster._check as check
+from azure.storage.blob import BlobServiceClient
+from dagster._core.pipes.client import PipesContextInjector, PipesParams
+from dagster_pipes import PipesContextData
+
+_CONTEXT_FILENAME = "context.json"
+
+
+class PipesAzureBlobStorageContextInjector(PipesContextInjector):
+    """A context injector that injects context by writing to a temporary AzureBlobStorage location.
+
+    Args:
+        container (str): The AzureBlobStorage container to write to.
+        client (azure.storage.blob.BlobServiceClient): An Azure Blob Storage client.
+        key_prefix (Optional[str]): An optional prefix to use for the Azure Blob Storage key. Defaults to a random
+            string.
+
+    """
+
+    def __init__(self, *, container: str, client: BlobServiceClient):
+        super().__init__()
+        self.bucket = check.str_param(container, "container")
+        self.client = client
+
+    @contextmanager
+    def inject_context(self, context: PipesContextData) -> Iterator[PipesParams]:  # pyright: ignore[reportIncompatibleMethodOverride]
+        key_prefix = "".join(random.choices(string.ascii_letters, k=30))
+        key = os.path.join(key_prefix, _CONTEXT_FILENAME)
+
+        with self.client.get_blob_client(self.bucket, key) as blob_client:
+            blob_client.upload_blob(json.dumps(context).encode("utf-8"))
+            yield {"bucket": self.bucket, "key": key}
+            blob_client.delete_blob()
+
+    def no_messages_debug_text(self) -> str:
+        return (
+            "Attempted to inject context via a temporary file in AzureBlobStorage. Expected"
+            " PipesAzureBlobStorageContextLoader to be explicitly passed to open_dagster_pipes in the external"
+            " process."
+        )

dagster_azure-0.28.0/dagster_azure/pipes/message_readers.py

@@ -0,0 +1,83 @@
+import random
+import string
+from collections.abc import Iterator, Sequence
+from contextlib import contextmanager
+from typing import Optional
+
+import dagster._check as check
+from azure.storage.blob import BlobServiceClient
+from dagster._core.pipes.client import PipesParams
+from dagster._core.pipes.utils import PipesBlobStoreMessageReader, PipesLogReader
+from dagster_pipes import PipesBlobStoreMessageWriter
+
+
+class PipesAzureBlobStorageMessageReader(PipesBlobStoreMessageReader):
+    """Message reader that reads messages by periodically reading message chunks from a specified AzureBlobStorage
+    container.
+
+    If `log_readers` is passed, this reader will also start the passed readers
+    when the first message is received from the external process.
+
+    Args:
+        interval (float): interval in seconds between attempts to download a chunk
+        container (str): The AzureBlobStorage container to read from.
+        client (azure.storage.blob.BlobServiceClient): An azure BlobServiceClient.
+        log_readers (Optional[Sequence[PipesLogReader]]): A set of log readers for logs on AzureBlobStorage.
+        include_stdio_in_messages (bool): Whether to send stdout/stderr to Dagster via Pipes messages. Defaults to False.
+    """
+
+    def __init__(
+        self,
+        *,
+        interval: float = 10,
+        container: str,
+        client: BlobServiceClient,
+        log_readers: Optional[Sequence[PipesLogReader]] = None,
+        include_stdio_in_messages: bool = False,
+    ):
+        super().__init__(
+            interval=interval,
+            log_readers=log_readers,
+        )
+        self.bucket = check.str_param(container, "container")
+        self.include_stdio_in_messages = check.bool_param(
+            include_stdio_in_messages, "include_stdio_in_messages"
+        )
+        self.client = client
+
+    @contextmanager
+    def get_params(self) -> Iterator[PipesParams]:
+        key_prefix = "".join(random.choices(string.ascii_letters, k=30))
+        yield {
+            "bucket": self.bucket,
+            "key_prefix": key_prefix,
+            PipesBlobStoreMessageWriter.INCLUDE_STDIO_IN_MESSAGES_KEY: self.include_stdio_in_messages,
+        }
+
+    def messages_are_readable(self, params: PipesParams) -> bool:
+        key_prefix = params.get("key_prefix")
+        if key_prefix is not None:
+            try:
+                with self.client.get_blob_client(
+                    self.bucket, f"{key_prefix}/1.json"
+                ) as blob_client:
+                    return blob_client.exists()
+            except Exception:
+                return False
+        else:
+            return False
+
+    def download_messages_chunk(self, index: int, params: PipesParams) -> Optional[str]:
+        key = f"{params['key_prefix']}/{index}.json"
+        try:
+            with self.client.get_blob_client(self.bucket, key) as blob_client:
+                return blob_client.download_blob().readall().decode("utf-8")
+        except Exception:
+            return None
+
+    def no_messages_debug_text(self) -> str:
+        return (
+            f"Attempted to read messages from Azure Blob Storage container "
+            f"{self.bucket}. Expected PipesAzureBlobStorageMessageWriter to be "
+            "explicitly passed to open_dagster_pipes in the external process."
+        )

dagster_azure-0.28.0/dagster_azure/version.py

@@ -0,0 +1 @@
+__version__ = "0.28.0"

{dagster_azure-0.27.15 → dagster_azure-0.28.0/dagster_azure.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dagster-azure
-Version: 0.27.15
+Version: 0.28.0
 Summary: Package for Azure-specific Dagster framework op and resource components.
 Home-page: https://github.com/dagster-io/dagster/tree/master/python_modules/libraries/dagster-azure
 Author: Dagster Labs
@@ -17,9 +17,10 @@ Requires-Python: >=3.9,<3.14
 License-File: LICENSE
 Requires-Dist: azure-core<2.0.0,>=1.7.0
 Requires-Dist: azure-identity<2.0.0,>=1.7.0
+Requires-Dist: azure-ai-ml<2.0.0,>=1.28.0
 Requires-Dist: azure-storage-blob<13.0.0,>=12.5.0
 Requires-Dist: azure-storage-file-datalake<13.0.0,>=12.5
-Requires-Dist: dagster==1.11.15
+Requires-Dist: dagster==1.12.0
 Dynamic: author
 Dynamic: author-email
 Dynamic: classifier

{dagster_azure-0.27.15 → dagster_azure-0.28.0}/dagster_azure.egg-info/SOURCES.txt

@@ -24,4 +24,9 @@ dagster_azure/blob/fake_blob_client.py
 dagster_azure/blob/resources.py
 dagster_azure/blob/utils.py
 dagster_azure/fakes/__init__.py
-dagster_azure/fakes/fake_adls2_resource.py
+dagster_azure/fakes/fake_adls2_resource.py
+dagster_azure/pipes/__init__.py
+dagster_azure/pipes/context_injectors.py
+dagster_azure/pipes/message_readers.py
+dagster_azure/pipes/clients/__init__.py
+dagster_azure/pipes/clients/azureml.py

{dagster_azure-0.27.15 → dagster_azure-0.28.0}/dagster_azure.egg-info/requires.txt

@@ -1,5 +1,6 @@
 azure-core<2.0.0,>=1.7.0
 azure-identity<2.0.0,>=1.7.0
+azure-ai-ml<2.0.0,>=1.28.0
 azure-storage-blob<13.0.0,>=12.5.0
 azure-storage-file-datalake<13.0.0,>=12.5
-dagster==1.11.15
+dagster==1.12.0

{dagster_azure-0.27.15 → dagster_azure-0.28.0}/setup.py

@@ -37,9 +37,10 @@ setup(
     install_requires=[
         "azure-core<2.0.0,>=1.7.0",
         "azure-identity<2.0.0,>=1.7.0",
+        "azure-ai-ml<2.0.0,>=1.28.0",
         "azure-storage-blob<13.0.0,>=12.5.0",
         "azure-storage-file-datalake<13.0.0,>=12.5",
-        "dagster==1.11.15",
+        "dagster==1.12.0",
     ],
     entry_points={"console_scripts": ["dagster-azure = dagster_azure.cli.cli:main"]},
     zip_safe=False,

dagster_azure-0.27.15/dagster_azure/version.py

@@ -1 +0,0 @@
-__version__ = "0.27.15"