wandb 0.13.10__py3-none-any.whl → 0.14.0__py3-none-any.whl

wandb/sdk/launch/runner/local_container.py

@@ -8,14 +8,17 @@ from typing import Any, Dict, List, Optional
 
 import wandb
 from wandb.sdk.launch.builder.abstract import AbstractBuilder
+from wandb.sdk.launch.environment.abstract import AbstractEnvironment
 
-from .._project_spec import LaunchProject, get_entry_point_command
-from ..builder.build import docker_image_exists, get_env_vars_dict, pull_docker_image
+from .._project_spec import LaunchProject, compute_command_args
+from ..builder.build import get_env_vars_dict
 from ..utils import (
     LOG_PREFIX,
     PROJECT_SYNCHRONOUS,
     _is_wandb_dev_uri,
     _is_wandb_local_uri,
+    docker_image_exists,
+    pull_docker_image,
     sanitize_wandb_api_key,
 )
 from .abstract import AbstractRun, AbstractRunner, Status
@@ -66,20 +69,24 @@ class LocalSubmittedRun(AbstractRun):
 class LocalContainerRunner(AbstractRunner):
     """Runner class, uses a project to create a LocallySubmittedRun."""
 
+    def __init__(
+        self,
+        api: wandb.apis.internal.Api,
+        backend_config: Dict[str, Any],
+        environment: AbstractEnvironment,
+    ) -> None:
+        super().__init__(api, backend_config)
+        self.environment = environment
+
     def run(
         self,
         launch_project: LaunchProject,
-        builder: AbstractBuilder,
-        registry_config: Dict[str, Any],
+        builder: Optional[AbstractBuilder],
     ) -> Optional[AbstractRun]:
         synchronous: bool = self.backend_config[PROJECT_SYNCHRONOUS]
         docker_args: Dict[str, Any] = launch_project.resource_args.get(
             "local-container", {}
         )
-        # TODO: leaving this here because of existing CLI command
-        # we should likely just tell users to specify the gpus arg directly
-        if launch_project.cuda:
-            docker_args["gpus"] = "all"
 
         if _is_wandb_local_uri(self._api.settings("base_url")):
             if sys.platform == "win32":
@@ -107,28 +114,39 @@ class LocalContainerRunner(AbstractRunner):
             image_uri = launch_project.image_name
             if not docker_image_exists(image_uri):
                 pull_docker_image(image_uri)
-            env_vars.pop("WANDB_RUN_ID")
-            # if they've given an override to the entrypoint
-            entry_cmd = get_entry_point_command(
-                entry_point, launch_project.override_args
-            )
+            entry_cmd = []
+            if entry_point is not None:
+                entry_cmd = entry_point.command
+            override_args = compute_command_args(launch_project.override_args)
             command_str = " ".join(
-                get_docker_command(image_uri, env_vars, entry_cmd, docker_args)
+                get_docker_command(
+                    image_uri,
+                    env_vars,
+                    entry_cmd=entry_cmd,
+                    docker_args=docker_args,
+                    additional_args=override_args,
+                )
             ).strip()
         else:
             assert entry_point is not None
-            repository: Optional[str] = registry_config.get("url")
+            _logger.info("Building docker image...")
+            assert builder is not None
             image_uri = builder.build_image(
                 launch_project,
-                repository,
                 entry_point,
             )
+            _logger.info(f"Docker image built with uri {image_uri}")
+            # entry_cmd and additional_args are empty here because
+            # if launch built the container they've been accounted
+            # in the dockerfile and env vars respectively
             command_str = " ".join(
-                get_docker_command(image_uri, env_vars, [""], docker_args)
+                get_docker_command(
+                    image_uri,
+                    env_vars,
+                    docker_args=docker_args,
+                )
             ).strip()
 
-        if not self.ack_run_queue_item(launch_project):
-            return None
         sanitized_cmd_str = sanitize_wandb_api_key(command_str)
         _msg = f"{LOG_PREFIX}Launching run in docker with command: {sanitized_cmd_str}"
         wandb.termlog(_msg)
@@ -170,10 +188,11 @@ def _run_entry_point(command: str, work_dir: Optional[str]) -> AbstractRun:
 def get_docker_command(
     image: str,
     env_vars: Dict[str, str],
-    entry_cmd: List[str],
+    entry_cmd: Optional[List[str]] = None,
     docker_args: Optional[Dict[str, Any]] = None,
+    additional_args: Optional[List[str]] = None,
 ) -> List[str]:
-    """Constructs the docker command using the image and docker args.
+    """Construct the docker command using the image and docker args.
 
     Arguments:
     image: a Docker image to be run
@@ -202,8 +221,13 @@ def get_docker_command(
             else:
                 cmd += [prefix, shlex.quote(str(value))]
 
+    if entry_cmd:
+        cmd += ["--entrypoint", entry_cmd[0]]
     cmd += [shlex.quote(image)]
-    cmd += entry_cmd
+    if entry_cmd and len(entry_cmd) > 1:
+        cmd += entry_cmd[1:]
+    if additional_args:
+        cmd += additional_args
     return cmd
 
 

wandb/sdk/launch/runner/local_process.py

@@ -3,13 +3,13 @@ import shlex
 from typing import Any, List, Optional
 
 import wandb
-from wandb.errors import LaunchError
 
 from .._project_spec import LaunchProject, get_entry_point_command
 from ..builder.build import get_env_vars_dict
 from ..utils import (
     LOG_PREFIX,
     PROJECT_SYNCHRONOUS,
+    LaunchError,
     _is_wandb_uri,
     download_wandb_python_deps,
     parse_wandb_uri,
@@ -81,9 +81,6 @@ class LocalProcessRunner(AbstractRunner):
         for env_key, env_value in env_vars.items():
             cmd += [f"{shlex.quote(env_key)}={shlex.quote(env_value)}"]
 
-        if not self.ack_run_queue_item(launch_project):
-            return None
-
         entry_cmd = get_entry_point_command(entry_point, launch_project.override_args)
         cmd += entry_cmd
 

wandb/sdk/launch/runner/{aws.py → sagemaker_runner.py}

@@ -1,23 +1,20 @@
-import configparser
+"""Implementation of the SageMakerRunner class."""
 import logging
-import os
-import subprocess
 import time
-from typing import Any, Dict, Optional, Tuple, cast
+from typing import Any, Dict, Optional, cast
 
 if False:
     import boto3  # type: ignore
 
 import wandb
-import wandb.docker as docker
 from wandb.apis.internal import Api
-from wandb.errors import LaunchError
 from wandb.sdk.launch.builder.abstract import AbstractBuilder
-from wandb.util import get_module
+from wandb.sdk.launch.environment.aws_environment import AwsEnvironment
+from wandb.sdk.launch.utils import LaunchError
 
 from .._project_spec import LaunchProject, get_entry_point_command
 from ..builder.build import get_env_vars_dict
-from ..utils import LOG_PREFIX, PROJECT_SYNCHRONOUS, run_shell, to_camel_case
+from ..utils import LOG_PREFIX, PROJECT_SYNCHRONOUS, to_camel_case
 from .abstract import AbstractRun, AbstractRunner, Status
 
 _logger = logging.getLogger(__name__)
@@ -69,32 +66,49 @@ class SagemakerSubmittedRun(AbstractRun):
         return self._status
 
 
-class AWSSagemakerRunner(AbstractRunner):
+class SageMakerRunner(AbstractRunner):
     """Runner class, uses a project to create a SagemakerSubmittedRun."""
 
+    def __init__(
+        self, api: Api, backend_config: Dict[str, Any], environment: AwsEnvironment
+    ) -> None:
+        """Initialize the SagemakerRunner.
+
+        Arguments:
+            api (Api): The API instance.
+            backend_config (Dict[str, Any]): The backend configuration.
+            environment (AwsEnvironment): The AWS environment.
+
+        Raises:
+            LaunchError: If the runner cannot be initialized.
+        """
+        super().__init__(api, backend_config)
+        self.environment = environment
+
     def run(
         self,
         launch_project: LaunchProject,
-        builder: AbstractBuilder,
-        registry_config: Dict[str, Any],
+        builder: Optional[AbstractBuilder],
     ) -> Optional[AbstractRun]:
-        _logger.info("using AWSSagemakerRunner")
+        """Run a project on Amazon Sagemaker.
 
-        boto3 = get_module(
-            "boto3",
-            "AWSSagemakerRunner requires boto3 to be installed,  install with pip install wandb[launch]",
-        )
-        botocore = get_module(
-            "botocore",
-            "AWSSagemakerRunner requires botocore to be installed,  install with pip install wandb[launch]",
-        )
+        Arguments:
+            launch_project (LaunchProject): The project to run.
+            builder (AbstractBuilder): The builder to use.
+
+        Returns:
+            Optional[AbstractRun]: The run instance.
+
+        Raises:
+            LaunchError: If the launch is unsuccessful.
+        """
+        _logger.info("using AWSSagemakerRunner")
 
         given_sagemaker_args = launch_project.resource_args.get("sagemaker")
         if given_sagemaker_args is None:
             raise LaunchError(
                 "No sagemaker args specified. Specify sagemaker args in resource_args"
             )
-        validate_sagemaker_requirements(given_sagemaker_args, registry_config)
 
         default_output_path = self.backend_config.get("runner", {}).get(
             "s3_output_path"
@@ -104,37 +118,22 @@ class AWSSagemakerRunner(AbstractRunner):
         ):
             default_output_path = f"s3://{default_output_path}"
 
-        region = get_region(given_sagemaker_args, registry_config.get("region"))
-        instance_role: bool = False
-        try:
-            client = boto3.client("sts")
-            instance_role = True
-            caller_id = client.get_caller_identity()
-
-        except botocore.exceptions.NoCredentialsError:
-            access_key, secret_key = get_aws_credentials(given_sagemaker_args)
-            client = boto3.client(
-                "sts", aws_access_key_id=access_key, aws_secret_access_key=secret_key
-            )
-            caller_id = client.get_caller_identity()
-
+        session = self.environment.get_session()
+        client = session.client("sts")
+        caller_id = client.get_caller_identity()
         account_id = caller_id["Account"]
+        _logger.info(f"Using account ID {account_id}")
         role_arn = get_role_arn(given_sagemaker_args, self.backend_config, account_id)
         entry_point = launch_project.get_single_entry_point()
+
+        # Create a sagemaker client to launch the job.
+        sagemaker_client = session.client("sagemaker")
+
         # if the user provided the image they want to use, use that, but warn it won't have swappable artifacts
         if (
             given_sagemaker_args.get("AlgorithmSpecification", {}).get("TrainingImage")
             is not None
         ):
-            if instance_role:
-                sagemaker_client = boto3.client("sagemaker", region_name=region)
-            else:
-                sagemaker_client = boto3.client(
-                    "sagemaker",
-                    region_name=region,
-                    aws_access_key_id=access_key,
-                    aws_secret_access_key=secret_key,
-                )
             sagemaker_args = build_sagemaker_args(
                 launch_project,
                 self._api,
@@ -152,57 +151,20 @@ class AWSSagemakerRunner(AbstractRunner):
                 run.wait()
             return run
 
-        _logger.info("Connecting to AWS ECR Client")
-        if instance_role:
-            ecr_client = boto3.client("ecr", region_name=region)
-        else:
-            ecr_client = boto3.client(
-                "ecr",
-                region_name=region,
-                aws_access_key_id=access_key,
-                aws_secret_access_key=secret_key,
-            )
-        repository = get_ecr_repository_url(
-            ecr_client, given_sagemaker_args, registry_config
-        )
-        # TODO: handle login credentials gracefully
-        login_credentials = registry_config.get("credentials")
-        if login_credentials is not None:
-            wandb.termwarn(
-                "Ignoring registry credentials for ECR, using those found on the system"
-            )
-
-        if builder.type != "kaniko":
-            _logger.info("Logging in to AWS ECR")
-            login_resp = aws_ecr_login(region, repository)
-            if login_resp is None or "Login Succeeded" not in login_resp:
-                raise LaunchError(f"Unable to login to ECR, response: {login_resp}")
-
         if launch_project.docker_image:
             image = launch_project.docker_image
         else:
             assert entry_point is not None
+            assert builder is not None
             # build our own image
+            _logger.info("Building docker image...")
             image = builder.build_image(
                 launch_project,
-                repository,
                 entry_point,
             )
-
-        if not self.ack_run_queue_item(launch_project):
-            return None
+            _logger.info(f"Docker image built with uri {image}")
 
         _logger.info("Connecting to sagemaker client")
-        if instance_role:
-            sagemaker_client = boto3.client("sagemaker", region_name=region)
-        else:
-            sagemaker_client = boto3.client(
-                "sagemaker",
-                region_name=region,
-                aws_access_key_id=access_key,
-                aws_secret_access_key=secret_key,
-            )
-
         command_args = get_entry_point_command(
             entry_point, launch_project.override_args
         )
@@ -225,29 +187,15 @@ class AWSSagemakerRunner(AbstractRunner):
         return run
 
 
-def aws_ecr_login(region: str, registry: str) -> Optional[str]:
-    pw_command = ["aws", "ecr", "get-login-password", "--region", region]
-    try:
-        pw = run_shell(pw_command)[0]
-    except subprocess.CalledProcessError:
-        raise LaunchError(
-            "Unable to get login password. Please ensure you have AWS credentials configured"
-        )
-    try:
-        docker_login_process = docker.login("AWS", pw, registry)
-    except Exception:
-        raise LaunchError(f"Failed to login to ECR {registry}")
-    return docker_login_process
-
-
 def merge_aws_tag_with_algorithm_specification(
     algorithm_specification: Optional[Dict[str, Any]], aws_tag: Optional[str]
 ) -> Dict[str, Any]:
-    """
-    AWS Sagemaker algorithms require a training image and an input mode.
-    If the user does not specify the specification themselves, define the spec
-    minimally using these two fields. Otherwise, if they specify the AlgorithmSpecification
-    set the training image if it is not set.
+    """Create an AWS AlgorithmSpecification.
+
+    AWS Sagemaker algorithms require a training image and an input mode. If the user
+    does not specify the specification themselves, define the spec minimally using these
+    two fields. Otherwise, if they specify the AlgorithmSpecification set the training
+    image if it is not set.
     """
     if algorithm_specification is None:
         return {
@@ -366,65 +314,10 @@ def launch_sagemaker_job(
     return run
 
 
-def get_region(
-    sagemaker_args: Dict[str, Any], registry_config_region: Optional[str] = None
-) -> str:
-    region = sagemaker_args.get("region")
-    if region is None:
-        region = registry_config_region
-    if region is None:
-        region = os.environ.get("AWS_DEFAULT_REGION")
-    if region is None and os.path.exists(os.path.expanduser("~/.aws/config")):
-        config = configparser.ConfigParser()
-        config.read(os.path.expanduser("~/.aws/config"))
-        section = sagemaker_args.get("profile") or "default"
-        try:
-            region = config.get(section, "region")
-        except (configparser.NoOptionError, configparser.NoSectionError):
-            raise LaunchError(
-                "Unable to detemine default region from ~/.aws/config. "
-                "Please specify region in resource args or specify config "
-                "section as 'profile'"
-            )
-
-    if region is None:
-        raise LaunchError(
-            "AWS region not specified and ~/.aws/config not found. Configure AWS"
-        )
-    assert isinstance(region, str)
-    return region
-
-
-def get_aws_credentials(sagemaker_args: Dict[str, Any]) -> Tuple[str, str]:
-    access_key = os.environ.get("AWS_ACCESS_KEY_ID")
-    secret_key = os.environ.get("AWS_SECRET_ACCESS_KEY")
-    if (
-        access_key is None
-        or secret_key is None
-        and os.path.exists(os.path.expanduser("~/.aws/credentials"))
-    ):
-        profile = sagemaker_args.get("profile") or "default"
-        config = configparser.ConfigParser()
-        config.read(os.path.expanduser("~/.aws/credentials"))
-        try:
-            access_key = config.get(profile, "aws_access_key_id")
-            secret_key = config.get(profile, "aws_secret_access_key")
-        except (configparser.NoOptionError, configparser.NoSectionError):
-            raise LaunchError(
-                "Unable to get aws credentials from ~/.aws/credentials. "
-                "Please set aws credentials in environments variables, or "
-                "check your credentials in ~/.aws/credentials. Use resource "
-                "args to specify the profile using 'profile'"
-            )
-
-    if access_key is None or secret_key is None:
-        raise LaunchError("AWS credentials not found")
-    return access_key, secret_key
-
-
 def get_role_arn(
     sagemaker_args: Dict[str, Any], backend_config: Dict[str, Any], account_id: str
 ) -> str:
+    """Get the role arn from the sagemaker args or the backend config."""
     role_arn = sagemaker_args.get("RoleArn") or sagemaker_args.get("role_arn")
     if role_arn is None:
         role_arn = backend_config.get("runner", {}).get("role_arn")
@@ -437,55 +330,3 @@ def get_role_arn(
         return role_arn
 
     return f"arn:aws:iam::{account_id}:role/{role_arn}"
-
-
-def validate_sagemaker_requirements(
-    given_sagemaker_args: Dict[str, Any], registry_config: Dict[str, Any]
-) -> None:
-    if (
-        given_sagemaker_args.get(
-            "EcrRepoName", given_sagemaker_args.get("ecr_repo_name")
-        )
-        is None
-        and registry_config.get("url") is None
-    ):
-        raise LaunchError(
-            "AWS sagemaker requires an ECR Repository to push the container to "
-            "set this by adding a `EcrRepoName` key to the sagemaker"
-            "field of resource_args or through the url key in the registry section "
-            "of the launch agent config."
-        )
-
-    if registry_config.get("ecr-repo-provider", "aws").lower() != "aws":
-        raise LaunchError(
-            "Sagemaker jobs requires an AWS ECR Repo to push the container to"
-        )
-
-
-def get_ecr_repository_url(
-    ecr_client: "boto3.Client",
-    given_sagemaker_args: Dict[str, Any],
-    registry_config: Dict[str, Any],
-) -> str:
-    token = ecr_client.get_authorization_token()
-    ecr_repo_name = given_sagemaker_args.get(
-        "EcrRepoName", given_sagemaker_args.get("ecr_repo_name")
-    )
-    if ecr_repo_name:
-        if not isinstance(ecr_repo_name, str):
-            raise LaunchError("EcrRepoName must be a string")
-        if not ecr_repo_name.startswith("arn:aws:ecr:"):
-            repository = cast(
-                str,
-                token["authorizationData"][0]["proxyEndpoint"].replace("https://", "")
-                + f"/{ecr_repo_name}",
-            )
-        else:
-            repository = ecr_repo_name
-    else:
-        repository = cast(str, registry_config.get("url", ""))
-    if not repository:
-        raise LaunchError(
-            "Must provide a repository url either through resource args or launch config file"
-        )
-    return repository