wandb 0.18.2__py3-none-musllinux_1_2_x86_64.whl

wandb/sdk/launch/agent/config.py

@@ -0,0 +1,296 @@
+"""Definition of the config object used by the Launch agent."""
+
+from enum import Enum
+from typing import List, Optional
+
+# ValidationError is imported for exception type checking purposes only.
+from pydantic import (  # type: ignore
+    BaseModel,
+    Field,
+    ValidationError,  # noqa: F401
+    root_validator,
+    validator,
+)
+
+import wandb
+from wandb.sdk.launch.utils import (
+    AZURE_BLOB_REGEX,
+    AZURE_CONTAINER_REGISTRY_URI_REGEX,
+    ELASTIC_CONTAINER_REGISTRY_URI_REGEX,
+    GCP_ARTIFACT_REGISTRY_URI_REGEX,
+    GCS_URI_RE,
+    S3_URI_RE,
+)
+
+__all__ = [
+    "ValidationError",
+    "AgentConfig",
+]
+
+
+class EnvironmentType(str, Enum):
+    """Enum of valid environment types."""
+
+    aws = "aws"
+    gcp = "gcp"
+    azure = "azure"
+
+
+class RegistryType(str, Enum):
+    """Enum of valid registry types."""
+
+    ecr = "ecr"
+    acr = "acr"
+    gcr = "gcr"
+
+
+class BuilderType(str, Enum):
+    """Enum of valid builder types."""
+
+    docker = "docker"
+    kaniko = "kaniko"
+    noop = "noop"
+
+
+class TargetPlatform(str, Enum):
+    """Enum of valid target platforms."""
+
+    linux_amd64 = "linux/amd64"
+    linux_arm64 = "linux/arm64"
+
+
+class RegistryConfig(BaseModel):
+    """Configuration for registry block.
+
+    Note that we don't forbid extra fields here because:
+    - We want to allow all fields supported by each registry
+    - We will perform validation on the registry object itself later
+    - Registry block is being deprecated in favor of destination field in builder
+    """
+
+    type: Optional[RegistryType] = Field(
+        None,
+        description="The type of registry to use.",
+    )
+    uri: Optional[str] = Field(
+        None,
+        description="The URI of the registry.",
+    )
+
+    @validator("uri")  # type: ignore
+    @classmethod
+    def validate_uri(cls, uri: str) -> str:
+        return validate_registry_uri(uri)
+
+
+class EnvironmentConfig(BaseModel):
+    """Configuration for the environment block."""
+
+    type: Optional[EnvironmentType] = Field(
+        None,
+        description="The type of environment to use.",
+    )
+    region: Optional[str] = Field(..., description="The region to use.")
+
+    class Config:
+        extra = "allow"
+
+    @root_validator(pre=True)  # type: ignore
+    @classmethod
+    def check_extra_fields(cls, values: dict) -> dict:
+        """Check for extra fields and print a warning."""
+        for key in values:
+            if key not in ["type", "region"]:
+                wandb.termwarn(
+                    f"Unrecognized field {key} in environment block. Please check your config file."
+                )
+        return values
+
+
+class BuilderConfig(BaseModel):
+    type: Optional[BuilderType] = Field(
+        None,
+        description="The type of builder to use.",
+    )
+    destination: Optional[str] = Field(
+        None,
+        description="The destination to use for the built image. If not provided, "
+        "the image will be pushed to the registry.",
+    )
+
+    platform: Optional[TargetPlatform] = Field(
+        None,
+        description="The platform to use for the built image. If not provided, "
+        "the platform will be detected automatically.",
+    )
+
+    build_context_store: Optional[str] = Field(
+        None,
+        description="The build context store to use. Required for kaniko builds.",
+        alias="build-context-store",
+    )
+    build_job_name: Optional[str] = Field(
+        "wandb-launch-container-build",
+        description="Name prefix of the build job.",
+        alias="build-job-name",
+    )
+    secret_name: Optional[str] = Field(
+        None,
+        description="The name of the secret to use for the build job.",
+        alias="secret-name",
+    )
+    secret_key: Optional[str] = Field(
+        None,
+        description="The key of the secret to use for the build job.",
+        alias="secret-key",
+    )
+    kaniko_image: Optional[str] = Field(
+        "gcr.io/kaniko-project/executor:latest",
+        description="The image to use for the kaniko executor.",
+        alias="kaniko-image",
+    )
+
+    @validator("build_context_store")  # type: ignore
+    @classmethod
+    def validate_build_context_store(
+        cls, build_context_store: Optional[str]
+    ) -> Optional[str]:
+        """Validate that the build context store is a valid container registry URI."""
+        if build_context_store is None:
+            return None
+        for regex in [
+            S3_URI_RE,
+            GCS_URI_RE,
+            AZURE_BLOB_REGEX,
+        ]:
+            if regex.match(build_context_store):
+                return build_context_store
+        raise ValueError(
+            "Invalid build context store. Build context store must be a URI for an "
+            "S3 bucket, GCS bucket, or Azure blob."
+        )
+
+    @root_validator(pre=True)  # type: ignore
+    @classmethod
+    def validate_docker(cls, values: dict) -> dict:
+        """Right now there are no required fields for docker builds."""
+        return values
+
+    @validator("destination")  # type: ignore
+    @classmethod
+    def validate_destination(cls, destination: Optional[str]) -> Optional[str]:
+        """Validate that the destination is a valid container registry URI."""
+        if destination is None:
+            return None
+        return validate_registry_uri(destination)
+
+
+class AgentConfig(BaseModel):
+    """Configuration for the Launch agent."""
+
+    queues: List[str] = Field(
+        default=[],
+        description="The queues to use for this agent.",
+    )
+    entity: Optional[str] = Field(
+        description="The W&B entity to use for this agent.",
+    )
+    max_jobs: Optional[int] = Field(
+        1,
+        description="The maximum number of jobs to run concurrently.",
+    )
+    max_schedulers: Optional[int] = Field(
+        1,
+        description="The maximum number of sweep schedulers to run concurrently.",
+    )
+    secure_mode: Optional[bool] = Field(
+        False,
+        description="Whether to use secure mode for this agent. If True, the "
+        "agent will reject runs that attempt to override the entrypoint or image.",
+    )
+    registry: Optional[RegistryConfig] = Field(
+        None,
+        description="The registry to use.",
+    )
+    environment: Optional[EnvironmentConfig] = Field(
+        None,
+        description="The environment to use.",
+    )
+    builder: Optional[BuilderConfig] = Field(
+        None,
+        description="The builder to use.",
+    )
+    verbosity: Optional[int] = Field(
+        0,
+        description="How verbose to print, 0 = default, 1 = verbose, 2 = very verbose",
+    )
+    stopped_run_timeout: Optional[int] = Field(
+        60,
+        description="How many seconds to wait after receiving the stop command before forcibly cancelling a run.",
+    )
+
+    class Config:
+        extra = "forbid"
+
+
+def validate_registry_uri(uri: str) -> str:
+    """Validate that the registry URI is a valid container registry URI.
+
+    The URI should resolve to an image name in a container registry. The recognized
+    formats are for ECR, ACR, and GCP Artifact Registry. If the URI does not match
+    any of these formats, a warning is printed indicating the registry type is not
+    recognized and the agent can't guarantee that images can be pushed.
+
+    If the format is recognized but does not resolve to an image name, an
+    error is raised. For example, if the URI is an ECR URI but does not include
+    an image name or includes a tag as well as an image name, an error is raised.
+    """
+    tag_msg = (
+        "Destination for built images may not include a tag, but the URI provided "
+        "includes the suffix '{tag}'. Please remove the tag and try again. The agent "
+        "will automatically tag each image with a unique hash of the source code."
+    )
+    if uri.startswith("https://"):
+        uri = uri[8:]
+
+    match = GCP_ARTIFACT_REGISTRY_URI_REGEX.match(uri)
+    if match:
+        if match.group("tag"):
+            raise ValueError(tag_msg.format(tag=match.group("tag")))
+        if not match.group("image_name"):
+            raise ValueError(
+                "An image name must be specified in the URI for a GCP Artifact Registry. "
+                "Please provide a uri with the format "
+                "'https://<region>-docker.pkg.dev/<project>/<repository>/<image>'."
+            )
+        return uri
+
+    match = AZURE_CONTAINER_REGISTRY_URI_REGEX.match(uri)
+    if match:
+        if match.group("tag"):
+            raise ValueError(tag_msg.format(tag=match.group("tag")))
+        if not match.group("repository"):
+            raise ValueError(
+                "A repository name must be specified in the URI for an "
+                "Azure Container Registry. Please provide a uri with the format "
+                "'https://<registry-name>.azurecr.io/<repository>'."
+            )
+        return uri
+
+    match = ELASTIC_CONTAINER_REGISTRY_URI_REGEX.match(uri)
+    if match:
+        if match.group("tag"):
+            raise ValueError(tag_msg.format(tag=match.group("tag")))
+        if not match.group("repository"):
+            raise ValueError(
+                "A repository name must be specified in the URI for an "
+                "Elastic Container Registry. Please provide a uri with the format "
+                "'https://<account-id>.dkr.ecr.<region>.amazonaws.com/<repository>'."
+            )
+        return uri
+
+    wandb.termwarn(
+        f"Unable to recognize registry type in URI {uri}. You are responsible "
+        "for ensuring the agent can push images to this registry."
+    )
+    return uri

wandb/sdk/launch/agent/job_status_tracker.py

@@ -0,0 +1,53 @@
+import logging
+from dataclasses import dataclass
+from typing import Optional
+
+from wandb.apis.internal import Api
+from wandb.errors import CommError
+from wandb.sdk.launch._project_spec import LaunchProject
+
+from ..runner.abstract import AbstractRun
+from ..utils import event_loop_thread_exec
+from .run_queue_item_file_saver import RunQueueItemFileSaver
+
+_logger = logging.getLogger(__name__)
+
+
+@dataclass
+class JobAndRunStatusTracker:
+    run_queue_item_id: str
+    queue: str
+    saver: RunQueueItemFileSaver
+    run_id: Optional[str] = None
+    project: Optional[str] = None
+    entity: Optional[str] = None
+    run: Optional[AbstractRun] = None
+    failed_to_start: bool = False
+    completed_status: Optional[str] = None
+    is_scheduler: bool = False
+    err_stage: str = "agent"
+
+    @property
+    def job_completed(self) -> bool:
+        return self.failed_to_start or self.completed_status is not None
+
+    def update_run_info(self, launch_project: LaunchProject) -> None:
+        self.run_id = launch_project.run_id
+        self.project = launch_project.target_project
+        self.entity = launch_project.target_entity
+
+    def set_err_stage(self, stage: str) -> None:
+        self.err_stage = stage
+
+    async def check_wandb_run_stopped(self, api: Api) -> bool:
+        assert (
+            self.run_id is not None
+            and self.project is not None
+            and self.entity is not None
+        ), "Job tracker does not contain run info. Update with run info before checking if run stopped"
+        check_stop = event_loop_thread_exec(api.api.check_stop_requested)
+        try:
+            return bool(await check_stop(self.project, self.entity, self.run_id))
+        except CommError as e:
+            _logger.error(f"CommError when checking if wandb run stopped: {e}")
+        return False

wandb/sdk/launch/agent/run_queue_item_file_saver.py

@@ -0,0 +1,45 @@
+"""Implementation of the run queue item file saver class."""
+
+import os
+import sys
+from typing import List, Optional
+
+import wandb
+
+if sys.version_info >= (3, 8):
+    from typing import Literal
+else:
+    from typing_extensions import Literal
+
+FileSubtypes = Literal["warning", "error"]
+
+
+class RunQueueItemFileSaver:
+    def __init__(
+        self,
+        agent_run: Optional["wandb.sdk.wandb_run.Run"],
+        run_queue_item_id: str,
+    ):
+        self.run_queue_item_id = run_queue_item_id
+        self.run = agent_run
+
+    def save_contents(
+        self, contents: str, fname: str, file_sub_type: FileSubtypes
+    ) -> Optional[List[str]]:
+        if not isinstance(self.run, wandb.sdk.wandb_run.Run):
+            wandb.termwarn("Not saving file contents because agent has no run")
+            return None
+        root_dir = self.run._settings.files_dir
+        saved_run_path = os.path.join(self.run_queue_item_id, file_sub_type, fname)
+        local_path = os.path.join(root_dir, saved_run_path)
+        os.makedirs(os.path.dirname(local_path), exist_ok=True)
+        with open(local_path, "w") as f:
+            f.write(contents)
+        res = self.run.save(local_path, base_path=root_dir, policy="now")
+        if isinstance(res, list):
+            return [saved_run_path]
+        else:
+            wandb.termwarn(
+                f"Failed to save files for run queue item: {self.run_queue_item_id}"
+            )
+            return None

wandb/sdk/launch/builder/abstract.py

@@ -0,0 +1,156 @@
+"""Abstract plugin class defining the interface needed to build container images for W&B Launch."""
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, Dict, Optional
+
+from wandb.sdk.launch.environment.abstract import AbstractEnvironment
+from wandb.sdk.launch.registry.abstract import AbstractRegistry
+
+from .._project_spec import EntryPoint, LaunchProject
+from ..registry.anon import AnonynmousRegistry
+from ..utils import (
+    AZURE_CONTAINER_REGISTRY_URI_REGEX,
+    ELASTIC_CONTAINER_REGISTRY_URI_REGEX,
+    GCP_ARTIFACT_REGISTRY_URI_REGEX,
+)
+
+if TYPE_CHECKING:
+    from wandb.sdk.launch.agent.job_status_tracker import JobAndRunStatusTracker
+
+
+class AbstractBuilder(ABC):
+    """Abstract plugin class defining the interface needed to build container images for W&B Launch."""
+
+    builder_type: str
+    environment: AbstractEnvironment
+    registry: AbstractRegistry
+    builder_config: Dict[str, Any]
+
+    @abstractmethod
+    def __init__(
+        self,
+        environment: AbstractEnvironment,
+        registry: AbstractRegistry,
+        verify: bool = True,
+    ) -> None:
+        """Initialize a builder.
+
+        Arguments:
+            builder_config: The builder config.
+            registry: The registry to use.
+            verify: Whether to verify the functionality of the builder.
+
+        Raises:
+            LaunchError: If the builder cannot be initialized or verified.
+        """
+        raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def from_config(
+        cls,
+        config: dict,
+        environment: AbstractEnvironment,
+        registry: AbstractRegistry,
+    ) -> "AbstractBuilder":
+        """Create a builder from a config dictionary.
+
+        Arguments:
+            config: The config dictionary.
+            environment: The environment to use.
+            registry: The registry to use.
+            verify: Whether to verify the functionality of the builder.
+            login: Whether to login to the registry immediately.
+
+        Returns:
+            The builder.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def build_image(
+        self,
+        launch_project: LaunchProject,
+        entrypoint: EntryPoint,
+        job_tracker: Optional["JobAndRunStatusTracker"] = None,
+    ) -> str:
+        """Build the image for the given project.
+
+        Arguments:
+            launch_project: The project to build.
+            build_ctx_path: The path to the build context.
+
+        Returns:
+            The image name.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def verify(self) -> None:
+        """Verify that the builder can be used to build images.
+
+        Raises:
+            LaunchError: If the builder cannot be used to build images.
+        """
+        raise NotImplementedError
+
+
+def registry_from_uri(uri: str) -> AbstractRegistry:
+    """Create a registry helper object from a uri.
+
+    This function parses the URI and determines which supported registry it
+    belongs to. It then creates a registry helper object for that registry.
+    The supported remote registry types are:
+    - Azure Container Registry
+    - Google Container Registry
+    - AWS Elastic Container Registry
+
+    The format of the URI is as follows:
+    - Azure Container Registry: <registry-name>.azurecr.io/<repo-name>/<image-name>
+    - Google Container Registry: <location>-docker.pkg.dev/<project-id>/<repo-name>/<image-name>
+    - AWS Elastic Container Registry: <account-id>.dkr.ecr.<region>.amazonaws.com/<repo-name>/<image-name>
+
+    Our classification of the registry is based on the domain name. For example,
+    if the uri contains `.azurecr.io`, we classify it as an Azure
+    Container Registry. If the uri contains `.dkr.ecr`, we classify
+    it as an AWS Elastic Container Registry. If the uri contains
+    `-docker.pkg.dev`, we classify it as a Google Artifact Registry.
+
+    This function will attempt to load the approriate cloud helpers for the
+
+    `https://` prefix is optional for all of the above.
+
+    Arguments:
+        uri: The uri to create a registry from.
+
+    Returns:
+        The registry.
+
+    Raises:
+        LaunchError: If the registry helper cannot be loaded for the given URI.
+    """
+    if uri.startswith("https://"):
+        uri = uri[len("https://") :]
+
+    if AZURE_CONTAINER_REGISTRY_URI_REGEX.match(uri) is not None:
+        from wandb.sdk.launch.registry.azure_container_registry import (
+            AzureContainerRegistry,
+        )
+
+        return AzureContainerRegistry(uri=uri)
+
+    elif GCP_ARTIFACT_REGISTRY_URI_REGEX.match(uri) is not None:
+        from wandb.sdk.launch.registry.google_artifact_registry import (
+            GoogleArtifactRegistry,
+        )
+
+        return GoogleArtifactRegistry(uri=uri)
+
+    elif ELASTIC_CONTAINER_REGISTRY_URI_REGEX.match(uri) is not None:
+        from wandb.sdk.launch.registry.elastic_container_registry import (
+            ElasticContainerRegistry,
+        )
+
+        return ElasticContainerRegistry(uri=uri)
+
+    return AnonynmousRegistry(uri=uri)