wandb 0.15.3__py3-none-any.whl → 0.15.5__py3-none-any.whl

wandb/sdk/launch/runner/local_container.py

@@ -1,12 +1,14 @@
 import logging
 import os
 import shlex
-import signal
 import subprocess
 import sys
+import threading
+import time
 from typing import Any, Dict, List, Optional
 
 import wandb
+from wandb.sdk.launch.agent.job_status_tracker import JobAndRunStatusTracker
 from wandb.sdk.launch.builder.abstract import AbstractBuilder
 from wandb.sdk.launch.environment.abstract import AbstractEnvironment
 
@@ -17,7 +19,6 @@ from ..utils import (
     PROJECT_SYNCHRONOUS,
     _is_wandb_dev_uri,
     _is_wandb_local_uri,
-    docker_image_exists,
     pull_docker_image,
     sanitize_wandb_api_key,
 )
@@ -29,36 +30,57 @@ _logger = logging.getLogger(__name__)
 class LocalSubmittedRun(AbstractRun):
     """Instance of ``AbstractRun`` corresponding to a subprocess launched to run an entry point command locally."""
 
-    def __init__(self, command_proc: "subprocess.Popen[bytes]") -> None:
+    def __init__(self) -> None:
         super().__init__()
-        self.command_proc = command_proc
+        self._command_proc: Optional[subprocess.Popen] = None
+        self._stdout: Optional[str] = None
+        self._terminate_flag: bool = False
+        self._thread: Optional[threading.Thread] = None
+
+    def set_command_proc(self, command_proc: subprocess.Popen) -> None:
+        self._command_proc = command_proc
+
+    def set_thread(self, thread: threading.Thread) -> None:
+        self._thread = thread
 
     @property
-    def id(self) -> str:
-        return str(self.command_proc.pid)
+    def id(self) -> Optional[str]:
+        if self._command_proc is None:
+            return None
+        return str(self._command_proc.pid)
 
     def wait(self) -> bool:
-        return self.command_proc.wait() == 0
+        assert self._thread is not None
+        # if command proc is not set
+        # wait for thread to set it
+        if self._command_proc is None:
+            while self._thread.is_alive():
+                time.sleep(5)
+                # command proc can be updated by another thread
+                if self._command_proc is not None:
+                    return self._command_proc.wait() == 0  # type: ignore
+            return False
+
+        return self._command_proc.wait() == 0
+
+    def get_logs(self) -> Optional[str]:
+        return self._stdout
 
     def cancel(self) -> None:
-        # Interrupt child process if it hasn't already exited
-        if self.command_proc.poll() is None:
-            # Kill the the process tree rooted at the child if it's the leader of its own process
-            # group, otherwise just kill the child
-            try:
-                if self.command_proc.pid == os.getpgid(self.command_proc.pid):
-                    os.killpg(self.command_proc.pid, signal.SIGTERM)
-                else:
-                    self.command_proc.terminate()
-            except OSError:
-                # The child process may have exited before we attempted to terminate it, so we
-                # ignore OSErrors raised during child process termination
-                _msg = f"{LOG_PREFIX}Failed to terminate child process PID {self.command_proc.pid}"
-                _logger.debug(_msg)
-            self.command_proc.wait()
+        # thread is set immediately after starting, should always exist
+        assert self._thread is not None
+
+        # cancel called before the thread subprocess has started
+        # indicates to thread to not start command proc if not already started
+        self._terminate_flag = True
 
     def get_status(self) -> Status:
-        exit_code = self.command_proc.poll()
+        assert self._thread is not None, "Failed to get status, self._thread = None"
+        if self._command_proc is None:
+            if self._thread.is_alive():
+                return Status("running")
+            return Status("stopped")
+        exit_code = self._command_proc.poll()
         if exit_code is None:
             return Status("running")
         if exit_code == 0:
@@ -78,12 +100,7 @@ class LocalContainerRunner(AbstractRunner):
         super().__init__(api, backend_config)
         self.environment = environment
 
-    def run(
-        self,
-        launch_project: LaunchProject,
-        builder: Optional[AbstractBuilder],
-    ) -> Optional[AbstractRun]:
-        synchronous: bool = self.backend_config[PROJECT_SYNCHRONOUS]
+    def _populate_docker_args(self, launch_project: LaunchProject) -> Dict[str, Any]:
         docker_args: Dict[str, Any] = launch_project.resource_args.get(
             "local-container", {}
         )
@@ -96,6 +113,16 @@ class LocalContainerRunner(AbstractRunner):
             if sys.platform == "linux" or sys.platform == "linux2":
                 docker_args["add-host"] = "host.docker.internal:host-gateway"
 
+        return docker_args
+
+    def run(
+        self,
+        launch_project: LaunchProject,
+        builder: Optional[AbstractBuilder],
+        job_tracker: Optional[JobAndRunStatusTracker] = None,
+    ) -> Optional[AbstractRun]:
+        docker_args = self._populate_docker_args(launch_project)
+        synchronous: bool = self.backend_config[PROJECT_SYNCHRONOUS]
         entry_point = launch_project.get_single_entry_point()
         env_vars = get_env_vars_dict(launch_project, self._api)
 
@@ -107,13 +134,12 @@ class LocalContainerRunner(AbstractRunner):
             _, _, port = self._api.settings("base_url").split(":")
             env_vars["WANDB_BASE_URL"] = f"http://host.docker.internal:{port}"
         elif _is_wandb_dev_uri(self._api.settings("base_url")):
-            env_vars["WANDB_BASE_URL"] = "http://host.docker.internal:9002"
+            env_vars["WANDB_BASE_URL"] = "http://host.docker.internal:9001"
 
         if launch_project.docker_image:
             # user has provided their own docker image
             image_uri = launch_project.image_name
-            if not docker_image_exists(image_uri):
-                pull_docker_image(image_uri)
+            pull_docker_image(image_uri)
             entry_cmd = []
             if entry_point is not None:
                 entry_cmd = entry_point.command
@@ -130,10 +156,7 @@ class LocalContainerRunner(AbstractRunner):
             assert entry_point is not None
             _logger.info("Building docker image...")
             assert builder is not None
-            image_uri = builder.build_image(
-                launch_project,
-                entry_point,
-            )
+            image_uri = builder.build_image(launch_project, entry_point, job_tracker)
             _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
@@ -145,7 +168,7 @@ class LocalContainerRunner(AbstractRunner):
                     docker_args=docker_args,
                 )
             ).strip()
-
+        launch_project.fill_macros(image_uri)
         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)
@@ -168,20 +191,49 @@ def _run_entry_point(command: str, work_dir: Optional[str]) -> AbstractRun:
     if work_dir is None:
         work_dir = os.getcwd()
     env = os.environ.copy()
-    if os.name == "nt":
-        # we are running on windows
-        process = subprocess.Popen(
-            ["cmd", "/c", command], close_fds=True, cwd=work_dir, env=env
-        )
-    else:
-        process = subprocess.Popen(
-            ["bash", "-c", command],
-            close_fds=True,
-            cwd=work_dir,
-            env=env,
-        )
-
-    return LocalSubmittedRun(process)
+    run = LocalSubmittedRun()
+    thread = threading.Thread(
+        target=_thread_process_runner,
+        args=(run, ["bash", "-c", command], work_dir, env),
+    )
+    run.set_thread(thread)
+    thread.start()
+    return run
+
+
+def _thread_process_runner(
+    run: LocalSubmittedRun, args: List[str], work_dir: str, env: Dict[str, str]
+) -> None:
+    # cancel was called before we started the subprocess
+    if run._terminate_flag:
+        return
+    process = subprocess.Popen(
+        args,
+        close_fds=True,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        universal_newlines=True,
+        bufsize=1,
+        cwd=work_dir,
+        env=env,
+    )
+    run.set_command_proc(process)
+    run._stdout = ""
+    while True:
+        # the agent thread could set the terminate flag
+        if run._terminate_flag:
+            process.terminate()  # type: ignore
+        chunk = os.read(process.stdout.fileno(), 4096)  # type: ignore
+        if not chunk:
+            break
+        index = chunk.find(b"\r")
+        decoded_chunk = chunk.decode()
+        if index != -1:
+            run._stdout += decoded_chunk
+            print(chunk.decode(), end="")
+        else:
+            run._stdout += decoded_chunk + "\r"
+            print(chunk.decode(), end="\r")
 
 
 def get_docker_command(

wandb/sdk/launch/runner/local_process.py

@@ -6,10 +6,10 @@ import wandb
 
 from .._project_spec import LaunchProject, get_entry_point_command
 from ..builder.build import get_env_vars_dict
+from ..errors import LaunchError
 from ..utils import (
     LOG_PREFIX,
     PROJECT_SYNCHRONOUS,
-    LaunchError,
     _is_wandb_uri,
     download_wandb_python_deps,
     parse_wandb_uri,

wandb/sdk/launch/runner/sagemaker_runner.py

@@ -8,9 +8,10 @@ if False:
 
 import wandb
 from wandb.apis.internal import Api
+from wandb.sdk.launch.agent.job_status_tracker import JobAndRunStatusTracker
 from wandb.sdk.launch.builder.abstract import AbstractBuilder
 from wandb.sdk.launch.environment.aws_environment import AwsEnvironment
-from wandb.sdk.launch.utils import LaunchError
+from wandb.sdk.launch.errors import LaunchError
 
 from .._project_spec import LaunchProject, get_entry_point_command
 from ..builder.build import get_env_vars_dict
@@ -23,9 +24,15 @@ _logger = logging.getLogger(__name__)
 class SagemakerSubmittedRun(AbstractRun):
     """Instance of ``AbstractRun`` corresponding to a subprocess launched to run an entry point command on aws sagemaker."""
 
-    def __init__(self, training_job_name: str, client: "boto3.Client") -> None:
+    def __init__(
+        self,
+        training_job_name: str,
+        client: "boto3.Client",
+        log_client: Optional["boto3.Client"] = None,
+    ) -> None:
         super().__init__()
         self.client = client
+        self.log_client = log_client
         self.training_job_name = training_job_name
         self._status = Status("running")
 
@@ -33,6 +40,38 @@ class SagemakerSubmittedRun(AbstractRun):
     def id(self) -> str:
         return f"sagemaker-{self.training_job_name}"
 
+    def get_logs(self) -> Optional[str]:
+        if self.log_client is None:
+            return None
+        try:
+            describe_res = self.log_client.describe_log_streams(
+                logGroupName="/aws/sagemaker/TrainingJobs",
+                logStreamNamePrefix=self.training_job_name,
+            )
+            if len(describe_res["logStreams"]) == 0:
+                wandb.termwarn(
+                    f"Failed to get logs for training job: {self.training_job_name}"
+                )
+                return None
+            log_name = describe_res["logStreams"][0]["logStreamName"]
+            res = self.log_client.get_log_events(
+                logGroupName="/aws/sagemaker/TrainingJobs",
+                logStreamName=log_name,
+            )
+            return "\n".join(
+                [f'{event["timestamp"]}:{event["message"]}' for event in res["events"]]
+            )
+        except self.log_client.exceptions.ResourceNotFoundException:
+            wandb.termwarn(
+                f"Failed to get logs for training job: {self.training_job_name}"
+            )
+            return None
+        except Exception as e:
+            wandb.termwarn(
+                f"Failed to handle logs for training job: {self.training_job_name} with error {str(e)}"
+            )
+            return None
+
     def wait(self) -> bool:
         while True:
             status_state = self.get_status().state
@@ -89,6 +128,7 @@ class SageMakerRunner(AbstractRunner):
         self,
         launch_project: LaunchProject,
         builder: Optional[AbstractBuilder],
+        job_tracker: Optional[JobAndRunStatusTracker] = None,
     ) -> Optional[AbstractRun]:
         """Run a project on Amazon Sagemaker.
 
@@ -128,6 +168,13 @@ class SageMakerRunner(AbstractRunner):
 
         # Create a sagemaker client to launch the job.
         sagemaker_client = session.client("sagemaker")
+        log_client = None
+        try:
+            log_client = session.client("logs")
+        except Exception as e:
+            wandb.termwarn(
+                f"Failed to connect to cloudwatch logs with error {str(e)}, logs will not be available"
+            )
 
         # if the user provided the image they want to use, use that, but warn it won't have swappable artifacts
         if (
@@ -146,7 +193,9 @@ class SageMakerRunner(AbstractRunner):
             _logger.info(
                 f"Launching sagemaker job on user supplied image with args: {sagemaker_args}"
             )
-            run = launch_sagemaker_job(launch_project, sagemaker_args, sagemaker_client)
+            run = launch_sagemaker_job(
+                launch_project, sagemaker_args, sagemaker_client, log_client
+            )
             if self.backend_config[PROJECT_SYNCHRONOUS]:
                 run.wait()
             return run
@@ -158,12 +207,10 @@ class SageMakerRunner(AbstractRunner):
             assert builder is not None
             # build our own image
             _logger.info("Building docker image...")
-            image = builder.build_image(
-                launch_project,
-                entry_point,
-            )
+            image = builder.build_image(launch_project, entry_point, job_tracker)
             _logger.info(f"Docker image built with uri {image}")
 
+        launch_project.fill_macros(image)
         _logger.info("Connecting to sagemaker client")
         command_args = get_entry_point_command(
             entry_point, launch_project.override_args
@@ -181,7 +228,9 @@ class SageMakerRunner(AbstractRunner):
             launch_project, self._api, role_arn, image, default_output_path
         )
         _logger.info(f"Launching sagemaker job with args: {sagemaker_args}")
-        run = launch_sagemaker_job(launch_project, sagemaker_args, sagemaker_client)
+        run = launch_sagemaker_job(
+            launch_project, sagemaker_args, sagemaker_client, log_client
+        )
         if self.backend_config[PROJECT_SYNCHRONOUS]:
             run.wait()
         return run
@@ -296,14 +345,15 @@ def launch_sagemaker_job(
     launch_project: LaunchProject,
     sagemaker_args: Dict[str, Any],
     sagemaker_client: "boto3.Client",
+    log_client: Optional["boto3.Client"] = None,
 ) -> SagemakerSubmittedRun:
     training_job_name = sagemaker_args.get("TrainingJobName") or launch_project.run_id
     resp = sagemaker_client.create_training_job(**sagemaker_args)
 
     if resp.get("TrainingJobArn") is None:
-        raise LaunchError("Unable to create training job")
+        raise LaunchError("Failed to create training job when submitting to SageMaker")
 
-    run = SagemakerSubmittedRun(training_job_name, sagemaker_client)
+    run = SagemakerSubmittedRun(training_job_name, sagemaker_client, log_client)
     wandb.termlog(
         f"{LOG_PREFIX}Run job submitted with arn: {resp.get('TrainingJobArn')}"
     )

wandb/sdk/launch/runner/vertex_runner.py

@@ -14,10 +14,12 @@ from wandb.apis.internal import Api
 from wandb.util import get_module
 
 from .._project_spec import LaunchProject, get_entry_point_command
+from ..agent.job_status_tracker import JobAndRunStatusTracker
 from ..builder.abstract import AbstractBuilder
 from ..builder.build import get_env_vars_dict
 from ..environment.gcp_environment import GcpEnvironment
-from ..utils import LOG_PREFIX, PROJECT_SYNCHRONOUS, LaunchError, run_shell
+from ..errors import LaunchError
+from ..utils import LOG_PREFIX, PROJECT_SYNCHRONOUS, run_shell
 from .abstract import AbstractRun, AbstractRunner, Status
 
 GCP_CONSOLE_URI = "https://console.cloud.google.com"
@@ -34,6 +36,10 @@ class VertexSubmittedRun(AbstractRun):
         # numeric ID of the custom training job
         return self._job.name  # type: ignore
 
+    def get_logs(self) -> Optional[str]:
+        # TODO: implement
+        return None
+
     @property
     def name(self) -> str:
         return self._job.display_name  # type: ignore
@@ -88,6 +94,7 @@ class VertexRunner(AbstractRunner):
         self,
         launch_project: LaunchProject,
         builder: Optional[AbstractBuilder],
+        job_tracker: Optional[JobAndRunStatusTracker] = None,
     ) -> Optional[AbstractRun]:
         """Run a Vertex job."""
         aiplatform = get_module(  # noqa: F811
@@ -133,11 +140,9 @@ class VertexRunner(AbstractRunner):
         else:
             assert entry_point is not None
             assert builder is not None
-            image_uri = builder.build_image(
-                launch_project,
-                entry_point,
-            )
+            image_uri = builder.build_image(launch_project, entry_point, job_tracker)
 
+        launch_project.fill_macros(image_uri)
         # TODO: how to handle this?
         entry_cmd = get_entry_point_command(entry_point, launch_project.override_args)
 

wandb/sdk/launch/sweeps/__init__.py

@@ -3,8 +3,6 @@ from typing import Any, Callable, Dict
 
 log = logging.getLogger(__name__)
 
-SCHEDULER_URI = "placeholder-uri-scheduler"
-
 
 class SchedulerError(Exception):
     """Raised when a known error occurs with wandb sweep scheduler."""
@@ -19,20 +17,20 @@ def _import_sweep_scheduler() -> Any:
 
 
 _WANDB_SCHEDULERS: Dict[str, Callable] = {
-    "sweep": _import_sweep_scheduler,
+    "wandb": _import_sweep_scheduler,
 }
 
 
-def load_scheduler(scheduler_name: str) -> Any:
-    scheduler_name = scheduler_name.lower()
-    if scheduler_name not in _WANDB_SCHEDULERS:
+def load_scheduler(scheduler_type: str) -> Any:
+    scheduler_type = scheduler_type.lower()
+    if scheduler_type not in _WANDB_SCHEDULERS:
         raise SchedulerError(
             f"The `scheduler_name` argument must be one of "
-            f"{list(_WANDB_SCHEDULERS.keys())}, got: {scheduler_name}"
+            f"{list(_WANDB_SCHEDULERS.keys())}, got: {scheduler_type}"
         )
 
-    log.warn(f"Loading dependencies for Scheduler of type: {scheduler_name}")
-    import_func = _WANDB_SCHEDULERS[scheduler_name]
+    log.warn(f"Loading dependencies for Scheduler of type: {scheduler_type}")
+    import_func = _WANDB_SCHEDULERS[scheduler_type]
     return import_func()