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

wandb/sdk/internal/tb_watcher.py

@@ -58,12 +58,16 @@ def _link_and_save_file(
     interface.publish_files(dict(files=[(GlobStr(glob.escape(file_name)), "live")]))
 
 
-def is_tfevents_file_created_by(path: str, hostname: str, start_time: float) -> bool:
-    """Check if a path is a tfevents file created by hostname.
+def is_tfevents_file_created_by(
+    path: str, hostname: Optional[str], start_time: Optional[float]
+) -> bool:
+    """Check if a path is a tfevents file.
+
+    Optionally checks that it was created by [hostname] after [start_time].
 
     tensorboard tfevents filename format:
         https://github.com/tensorflow/tensorboard/blob/f3f26b46981da5bd46a5bb93fcf02d9eb7608bc1/tensorboard/summary/writer/event_file_writer.py#L81
-    tensorflow tfevents fielname format:
+    tensorflow tfevents filename format:
         https://github.com/tensorflow/tensorflow/blob/8f597046dc30c14b5413813d02c0e0aed399c177/tensorflow/core/util/events_writer.cc#L68
     """
     if not path:
@@ -77,23 +81,27 @@ def is_tfevents_file_created_by(path: str, hostname: str, start_time: float) ->
     except ValueError:
         return False
     # check the hostname, which may have dots
-    for i, part in enumerate(hostname.split(".")):
+    if hostname is not None:
+        for i, part in enumerate(hostname.split(".")):
+            try:
+                fname_component_part = fname_components[tfevents_idx + 2 + i]
+            except IndexError:
+                return False
+            if part != fname_component_part:
+                return False
+    if start_time is not None:
         try:
-            fname_component_part = fname_components[tfevents_idx + 2 + i]
-        except IndexError:
+            created_time = int(fname_components[tfevents_idx + 1])
+        except (ValueError, IndexError):
             return False
-        if part != fname_component_part:
+        # Ensure that the file is newer then our start time, and that it was
+        # created from the same hostname.
+        # TODO: we should also check the PID (also contained in the tfevents
+        #     filename). Can we assume that our parent pid is the user process
+        #     that wrote these files?
+        if created_time < int(start_time):
             return False
-    try:
-        created_time = int(fname_components[tfevents_idx + 1])
-    except (ValueError, IndexError):
-        return False
-    # Ensure that the file is newer then our start time, and that it was
-    # created from the same hostname.
-    # TODO: we should also check the PID (also contained in the tfevents
-    #     filename). Can we assume that our parent pid is the user process
-    #     that wrote these files?
-    return created_time >= int(start_time)
+    return True
 
 
 class TBWatcher:
@@ -136,6 +144,7 @@ class TBWatcher:
             # Note that we strip '/' instead of os.sep, because elsewhere we've
             # converted paths to forward slash.
             namespace = logdir.replace(filename, "").replace(rootdir, "").strip("/")
+
             # TODO: revisit this heuristic, it exists because we don't know the
             # root log directory until more than one tfevents file is written to
             if len(dirs) == 1 and namespace not in ["train", "validation"]:
@@ -217,12 +226,13 @@ class TBDirWatcher:
         """Check if a path has been modified since launch and contains tfevents."""
         if not path:
             raise ValueError("Path must be a nonempty string")
-        if self._force:
-            return True
         path = self.tf_compat.tf.compat.as_str_any(path)
-        return is_tfevents_file_created_by(
-            path, self._hostname, self._tbwatcher._settings._start_time
-        )
+        if self._force:
+            return is_tfevents_file_created_by(path, None, None)
+        else:
+            return is_tfevents_file_created_by(
+                path, self._hostname, self._tbwatcher._settings._start_time
+            )
 
     def _loader(
         self, save: bool = True, namespace: Optional[str] = None

wandb/sdk/internal/thread_local_settings.py

@@ -0,0 +1,18 @@
+import threading
+from typing import Dict, Optional
+
+
+# Context variable for setting API settings (api keys, etc.) for internal and public apis thread-locally
+# TODO: move this into actual settings
+class _ThreadLocalApiSettings(threading.local):
+    api_key: Optional[str]
+    cookies: Optional[Dict]
+    headers: Optional[Dict]
+
+    def __init__(self) -> None:
+        self.api_key = None
+        self.cookies = None
+        self.headers = None
+
+
+_thread_local_api_settings: _ThreadLocalApiSettings = _ThreadLocalApiSettings()

wandb/sdk/launch/_project_spec.py

@@ -7,17 +7,20 @@ import json
 import logging
 import os
 import tempfile
-from typing import Any, Dict, List, Optional
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
 import wandb
 import wandb.docker as docker
 from wandb.apis.internal import Api
-from wandb.apis.public import Artifact as PublicArtifact
 from wandb.errors import CommError
+from wandb.sdk.launch import utils
 from wandb.sdk.lib.runid import generate_id
 
-from . import utils
-from .utils import LOG_PREFIX, LaunchError
+from .errors import LaunchError
+from .utils import LOG_PREFIX, recursive_macro_sub
+
+if TYPE_CHECKING:
+    from wandb.sdk.artifacts.artifact import Artifact
 
 _logger = logging.getLogger(__name__)
 
@@ -59,6 +62,7 @@ class LaunchProject:
         resource: str,
         resource_args: Dict[str, Any],
         run_id: Optional[str],
+        sweep_id: Optional[str] = None,
     ):
         if uri is not None and utils.is_bare_wandb_uri(uri):
             uri = api.settings("base_url") + uri
@@ -67,7 +71,7 @@ class LaunchProject:
         self.job = job
         if job is not None:
             wandb.termlog(f"{LOG_PREFIX}Launching job: {job}")
-        self._job_artifact: Optional[PublicArtifact] = None
+        self._job_artifact: Optional["Artifact"] = None
         self.api = api
         self.launch_spec = launch_spec
         self.target_entity = target_entity
@@ -78,11 +82,12 @@ class LaunchProject:
         # runner, so we need to pop the builder key out
         resource_args_build = resource_args.get(resource, {}).pop("builder", {})
         self.resource = resource
-        self.resource_args = resource_args
+        self.resource_args = resource_args.copy()
+        self.sweep_id = sweep_id
         self.python_version: Optional[str] = launch_spec.get("python_version")
-        self.cuda_base_image: Optional[str] = resource_args_build.get("cuda", {}).get(
-            "base_image"
-        )
+        self.accelerator_base_image: Optional[str] = resource_args_build.get(
+            "accelerator", {}
+        ).get("base_image") or resource_args_build.get("cuda", {}).get("base_image")
         self._base_image: Optional[str] = launch_spec.get("base_image")
         self.docker_image: Optional[str] = docker_config.get(
             "docker_image"
@@ -110,6 +115,9 @@ class LaunchProject:
             self.override_entrypoint = self.add_entry_point(
                 overrides.get("entry_point")  # type: ignore
             )
+        if overrides.get("sweep_id") is not None:
+            _logger.info("Adding override sweep id")
+            self.sweep_id = overrides["sweep_id"]
         if self.docker_image is not None:
             self.source = LaunchSource.DOCKER
             self.project_dir = None
@@ -172,6 +180,43 @@ class LaunchProject:
             assert self.job is not None
             return wandb.util.make_docker_image_name_safe(self.job.split(":")[0])
 
+    def fill_macros(self, image: str) -> None:
+        """Substitute values for macros in resource arguments.
+
+        Certain macros can be used in resource args. These macros allow the
+        user to set resource args dynamically in the context of the
+        run being launched. The macros are given in the ${macro} format. The
+        following macros are currently supported:
+
+        ${project_name} - the name of the project the run is being launched to.
+        ${entity_name} - the owner of the project the run being launched to.
+        ${run_id} - the id of the run being launched.
+        ${run_name} - the name of the run that is launching.
+        ${image_uri} - the URI of the container image for this run.
+
+        Additionally, you may use ${<ENV-VAR-NAME>} to refer to the value of any
+        environment variables that you plan to set in the environment of any
+        agents that will receive these resource args.
+
+        Calling this method will overwrite the contents of self.resource_args
+        with the substituted values.
+
+        Args:
+            image (str): The image name to fill in for ${wandb-image}.
+
+        Returns:
+            None
+        """
+        update_dict = {
+            "project_name": self.target_project,
+            "entity_name": self.target_entity,
+            "run_id": self.run_id,
+            "run_name": self.name,
+            "image_uri": image,
+        }
+        update_dict.update(os.environ)
+        self.resource_args = recursive_macro_sub(self.resource_args, update_dict)
+
     def build_required(self) -> bool:
         """Checks the source to see if a build is required."""
         # since the image tag for images built from jobs
@@ -416,6 +461,7 @@ def create_project_from_spec(launch_spec: Dict[str, Any], api: Api) -> LaunchPro
         launch_spec.get("resource", None),
         launch_spec.get("resource_args", {}),
         launch_spec.get("run_id", None),
+        launch_spec.get("sweep_id", {}),
     )
 
 
@@ -446,8 +492,8 @@ def fetch_and_validate_project(
         launch_project._fetch_project_local(internal_api=api)
 
     assert launch_project.project_dir is not None
-    # this prioritizes pip, and we don't support any cases where both are present
-    # conda projects when uploaded to wandb become pip projects via requirements.frozen.txt, wandb doesn't preserve conda envs
+    # this prioritizes pip, and we don't support any cases where both are present conda projects when uploaded to
+    # wandb become pip projects via requirements.frozen.txt, wandb doesn't preserve conda envs
     if os.path.exists(
         os.path.join(launch_project.project_dir, "requirements.txt")
     ) or os.path.exists(

wandb/sdk/launch/agent/agent.py

@@ -5,7 +5,6 @@ import pprint
 import threading
 import time
 import traceback
-from dataclasses import dataclass
 from multiprocessing import Event
 from multiprocessing.pool import ThreadPool
 from typing import Any, Dict, List, Optional, Union
@@ -13,22 +12,18 @@ from typing import Any, Dict, List, Optional, Union
 import wandb
 from wandb.apis.internal import Api
 from wandb.errors import CommError
-from wandb.sdk.launch._project_spec import LaunchProject
+from wandb.sdk.launch.launch_add import launch_add
 from wandb.sdk.launch.runner.local_container import LocalSubmittedRun
-from wandb.sdk.launch.sweeps import SCHEDULER_URI
+from wandb.sdk.launch.sweeps.scheduler import Scheduler
 from wandb.sdk.lib import runid
 
 from .. import loader
 from .._project_spec import create_project_from_spec, fetch_and_validate_project
 from ..builder.build import construct_builder_args
-from ..runner.abstract import AbstractRun
-from ..utils import (
-    LAUNCH_DEFAULT_PROJECT,
-    LOG_PREFIX,
-    PROJECT_SYNCHRONOUS,
-    LaunchDockerError,
-    LaunchError,
-)
+from ..errors import LaunchDockerError, LaunchError
+from ..utils import LAUNCH_DEFAULT_PROJECT, LOG_PREFIX, PROJECT_SYNCHRONOUS
+from .job_status_tracker import JobAndRunStatusTracker
+from .run_queue_item_file_saver import RunQueueItemFileSaver
 
 AGENT_POLLING_INTERVAL = 10
 ACTIVE_SWEEP_POLLING_INTERVAL = 1  # more frequent when we know we have jobs
@@ -37,30 +32,13 @@ AGENT_POLLING = "POLLING"
 AGENT_RUNNING = "RUNNING"
 AGENT_KILLED = "KILLED"
 
-MAX_THREADS = 64
-
-_logger = logging.getLogger(__name__)
+HIDDEN_AGENT_RUN_TYPE = "sweep-controller"
 
+MAX_THREADS = 64
 
-@dataclass
-class JobAndRunStatus:
-    run_queue_item_id: str
-    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
-
-    @property
-    def job_completed(self) -> bool:
-        return self.failed_to_start or self.completed_status is not None
+MAX_RESUME_COUNT = 5
 
-    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
+_logger = logging.getLogger(__name__)
 
 
 def _convert_access(access: str) -> str:
@@ -101,16 +79,21 @@ def _max_from_config(
     return max_from_config
 
 
-def _job_is_scheduler(run_spec: Dict[str, Any]) -> bool:
+def _is_scheduler_job(run_spec: Dict[str, Any]) -> bool:
     """Determine whether a job/runSpec is a sweep scheduler."""
     if not run_spec:
-        _logger.debug("Recieved runSpec in _job_is_scheduler that was empty")
+        _logger.debug("Recieved runSpec in _is_scheduler_job that was empty")
 
-    if run_spec.get("uri") != SCHEDULER_URI:
+    if run_spec.get("uri") != Scheduler.PLACEHOLDER_URI:
         return False
 
     if run_spec.get("resource") == "local-process":
-        # If a scheduler is a local-process (100%), also
+        # Any job pushed to a run queue that has a scheduler uri is
+        # allowed to use local-process
+        if run_spec.get("job"):
+            return True
+
+        # If a scheduler is local-process and run through CLI, also
         #    confirm command is in format: [wandb scheduler <sweep>]
         cmd = run_spec.get("overrides", {}).get("entry_point", [])
         if len(cmd) < 3:
@@ -137,7 +120,7 @@ class LaunchAgent:
         self._api = api
         self._base_url = self._api.settings().get("base_url")
         self._ticks = 0
-        self._jobs: Dict[int, JobAndRunStatus] = {}
+        self._jobs: Dict[int, JobAndRunStatusTracker] = {}
         self._jobs_lock = threading.Lock()
         self._jobs_event = Event()
         self._jobs_event.set()
@@ -169,15 +152,40 @@ class LaunchAgent:
             self.gorilla_supports_agents,
         )
         self._id = create_response["launchAgentId"]
-        self._name = ""  # hacky: want to display this to the user but we don't get it back from gql until polling starts. fix later
         if self._api.entity_is_team(self._entity):
             wandb.termwarn(
                 f"{LOG_PREFIX}Agent is running on team entity ({self._entity}). Members of this team will be able to run code on this device."
             )
 
-    def fail_run_queue_item(self, run_queue_item_id: str) -> None:
+        agent_response = self._api.get_launch_agent(
+            self._id, self.gorilla_supports_agents
+        )
+        self._name = agent_response["name"]
+        self._init_agent_run()
+
+    def fail_run_queue_item(
+        self,
+        run_queue_item_id: str,
+        message: str,
+        phase: str,
+        files: Optional[List[str]] = None,
+    ) -> None:
         if self._gorilla_supports_fail_run_queue_items:
-            self._api.fail_run_queue_item(run_queue_item_id)
+            self._api.fail_run_queue_item(run_queue_item_id, message, phase, files)
+
+    def _init_agent_run(self) -> None:
+        # TODO: has it been long enough that all backends support agents?
+        if self.gorilla_supports_agents:
+            settings = wandb.Settings(silent=True, disable_git=True)
+            self._wandb_run = wandb.init(
+                project=self._project,
+                entity=self._entity,
+                settings=settings,
+                id=self._name,
+                job_type=HIDDEN_AGENT_RUN_TYPE,
+            )
+        else:
+            self._wandb_run = None
 
     @property
     def thread_ids(self) -> List[int]:
@@ -253,24 +261,43 @@ class LaunchAgent:
         if not update_ret["success"]:
             wandb.termerror(f"{LOG_PREFIX}Failed to update agent status to {status}")
 
-    def finish_thread_id(self, thread_id: int) -> None:
+    def finish_thread_id(
+        self,
+        thread_id: int,
+        exception: Optional[Union[Exception, LaunchDockerError]] = None,
+    ) -> None:
         """Removes the job from our list for now."""
         job_and_run_status = self._jobs[thread_id]
-        if not job_and_run_status.run_id or not job_and_run_status.project:
-            self.fail_run_queue_item(job_and_run_status.run_queue_item_id)
-        elif job_and_run_status.entity != self._entity:
+        if (
+            job_and_run_status.entity is not None
+            and job_and_run_status.entity != self._entity
+        ):
             _logger.info(
                 "Skipping check for completed run status because run is on a different entity than agent"
             )
+        elif exception is not None:
+            tb_str = traceback.format_exception(
+                type(exception), value=exception, tb=exception.__traceback__
+            )
+            fnames = job_and_run_status.saver.save_contents(
+                "".join(tb_str), "error.log", "error"
+            )
+            self.fail_run_queue_item(
+                job_and_run_status.run_queue_item_id,
+                str(exception),
+                job_and_run_status.err_stage,
+                fnames,
+            )
         elif job_and_run_status.completed_status not in ["stopped", "failed"]:
             _logger.info(
                 "Skipping check for completed run status because run was successful"
             )
-        else:
+        elif job_and_run_status.run is not None:
             run_info = None
             # sweep runs exist but have no info before they are started
             # so run_info returned will be None
             # normal runs just throw a comm error
+            # TODO: make more clear
             try:
                 run_info = self._api.get_run_info(
                     self._entity, job_and_run_status.project, job_and_run_status.run_id
@@ -279,7 +306,22 @@ class LaunchAgent:
             except CommError:
                 pass
             if run_info is None:
-                self.fail_run_queue_item(job_and_run_status.run_queue_item_id)
+                _msg = "The submitted run was not successfully started"
+                fnames = None
+
+                logs = job_and_run_status.run.get_logs()
+                if logs:
+                    fnames = job_and_run_status.saver.save_contents(
+                        logs, "error.log", "error"
+                    )
+                self.fail_run_queue_item(
+                    job_and_run_status.run_queue_item_id, _msg, "run", fnames
+                )
+        else:
+            _logger.info("Finish thread id had no exception, ror run")
+            wandb._sentry.exception(
+                "launch agent called finish thread id on thread without run or exception"
+            )
 
         # TODO:  keep logs or something for the finished jobs
         with self._jobs_lock:
@@ -296,7 +338,9 @@ class LaunchAgent:
         if job.job_completed:
             self.finish_thread_id(thread_id)
 
-    def run_job(self, job: Dict[str, Any]) -> None:
+    def run_job(
+        self, job: Dict[str, Any], queue: str, file_saver: RunQueueItemFileSaver
+    ) -> None:
         """Set up project and run the job.
 
         Arguments:
@@ -322,6 +366,8 @@ class LaunchAgent:
                 job,
                 self.default_config,
                 self._api,
+                queue,
+                file_saver,
             ),
         )
 
@@ -367,7 +413,6 @@ class LaunchAgent:
                 agent_response = self._api.get_launch_agent(
                     self._id, self.gorilla_supports_agents
                 )
-                self._name = agent_response["name"]  # hack: first time we get name
                 if agent_response["stopPolling"]:
                     # shutdown process and all jobs if requested from ui
                     raise KeyboardInterrupt
@@ -376,7 +421,10 @@ class LaunchAgent:
                     for queue in self._queues:
                         job = self.pop_from_queue(queue)
                         if job:
-                            if _job_is_scheduler(job.get("runSpec")):
+                            file_saver = RunQueueItemFileSaver(
+                                self._wandb_run, job["runQueueItemId"]
+                            )
+                            if _is_scheduler_job(job.get("runSpec")):
                                 # If job is a scheduler, and we are already at the cap, ignore,
                                 #    don't ack, and it will be pushed back onto the queue in 1 min
                                 if self.num_running_schedulers >= self._max_schedulers:
@@ -388,13 +436,25 @@ class LaunchAgent:
                                     continue
 
                             try:
-                                self.run_job(job)
+                                self.run_job(job, queue, file_saver)
                             except Exception as e:
                                 wandb.termerror(
                                     f"{LOG_PREFIX}Error running job: {traceback.format_exc()}"
                                 )
                                 wandb._sentry.exception(e)
-                                self.fail_run_queue_item(job["runQueueItemId"])
+
+                                # always the first phase, because we only enter phase 2 within the thread
+                                files = file_saver.save_contents(
+                                    contents=traceback.format_exc(),
+                                    fname="error.log",
+                                    file_sub_type="error",
+                                )
+                                self.fail_run_queue_item(
+                                    run_queue_item_id=job["runQueueItemId"],
+                                    message=str(e),
+                                    phase="agent",
+                                    files=files,
+                                )
 
                 for thread_id in self.thread_ids:
                     self._update_finished(thread_id)
@@ -429,20 +489,27 @@ class LaunchAgent:
         job: Dict[str, Any],
         default_config: Dict[str, Any],
         api: Api,
+        queue: str,
+        file_saver: RunQueueItemFileSaver,
     ) -> None:
         thread_id = threading.current_thread().ident
         assert thread_id is not None
+        job_tracker = JobAndRunStatusTracker(job["runQueueItemId"], queue, file_saver)
+        with self._jobs_lock:
+            self._jobs[thread_id] = job_tracker
         try:
-            self._thread_run_job(launch_spec, job, default_config, api, thread_id)
+            self._thread_run_job(
+                launch_spec, job, default_config, api, queue, thread_id, job_tracker
+            )
         except LaunchDockerError as e:
             wandb.termerror(
                 f"{LOG_PREFIX}agent {self._name} encountered an issue while starting Docker, see above output for details."
             )
-            self.finish_thread_id(thread_id)
+            self.finish_thread_id(thread_id, e)
             wandb._sentry.exception(e)
         except Exception as e:
             wandb.termerror(f"{LOG_PREFIX}Error running job: {traceback.format_exc()}")
-            self.finish_thread_id(thread_id)
+            self.finish_thread_id(thread_id, e)
             wandb._sentry.exception(e)
 
     def _thread_run_job(
@@ -451,11 +518,10 @@ class LaunchAgent:
         job: Dict[str, Any],
         default_config: Dict[str, Any],
         api: Api,
+        queue: str,
         thread_id: int,
+        job_tracker: JobAndRunStatusTracker,
     ) -> None:
-        job_tracker = JobAndRunStatus(job["runQueueItemId"])
-        with self._jobs_lock:
-            self._jobs[thread_id] = job_tracker
         project = create_project_from_spec(launch_spec, api)
         job_tracker.update_run_info(project)
         _logger.info("Fetching and validating project...")
@@ -480,9 +546,8 @@ class LaunchAgent:
         backend = loader.runner_from_config(resource, api, backend_config, environment)
         _logger.info("Backend loaded...")
         api.ack_run_queue_item(job["runQueueItemId"], project.run_id)
-        run = backend.run(project, builder)
-
-        if _job_is_scheduler(launch_spec):
+        run = backend.run(project, builder, job_tracker)
+        if _is_scheduler_job(launch_spec):
             with self._jobs_lock:
                 self._jobs[thread_id].is_scheduler = True
             wandb.termlog(
@@ -497,15 +562,17 @@ class LaunchAgent:
         with self._jobs_lock:
             job_tracker.run = run
         while self._jobs_event.is_set():
-            if self._check_run_finished(job_tracker):
+            if self._check_run_finished(job_tracker, launch_spec):
                 return
             time.sleep(AGENT_POLLING_INTERVAL)
         # temp: for local, kill all jobs. we don't yet have good handling for different
         # types of runners in general
-        if isinstance(run, LocalSubmittedRun):
-            run.command_proc.kill()
+        if isinstance(run, LocalSubmittedRun) and run._command_proc is not None:
+            run._command_proc.kill()
 
-    def _check_run_finished(self, job_tracker: JobAndRunStatus) -> bool:
+    def _check_run_finished(
+        self, job_tracker: JobAndRunStatusTracker, launch_spec: Dict[str, Any]
+    ) -> bool:
         if job_tracker.completed_status:
             return True
 
@@ -522,13 +589,28 @@ class LaunchAgent:
         try:
             run = job_tracker.run
             status = run.get_status().state
-            if status in ["stopped", "failed", "finished"]:
+            if status in ["stopped", "failed", "finished", "preempted"]:
                 if job_tracker.is_scheduler:
                     wandb.termlog(f"{LOG_PREFIX}Scheduler finished with ID: {run.id}")
                 else:
                     wandb.termlog(f"{LOG_PREFIX}Job finished with ID: {run.id}")
                 with self._jobs_lock:
                     job_tracker.completed_status = status
+                if status == "preempted":
+                    config = launch_spec.copy()
+                    config["run_id"] = job_tracker.run_id
+                    config["_resume_count"] = config.get("_resume_count", 0) + 1
+                    if config["_resume_count"] > MAX_RESUME_COUNT:
+                        wandb.termlog(
+                            f"{LOG_PREFIX}Run {job_tracker.run_id} has already resumed {MAX_RESUME_COUNT} times."
+                        )
+                        return True
+                    wandb.termlog(f"{LOG_PREFIX}Requeueing run {job_tracker.run_id}.")
+                    launch_add(
+                        config=config,
+                        project_queue=self._project,
+                        queue_name=job_tracker.queue,
+                    )
                 return True
             return False
         except LaunchError as e:

wandb/sdk/launch/agent/job_status_tracker.py

@@ -0,0 +1,34 @@
+from dataclasses import dataclass
+from typing import Optional
+
+from wandb.sdk.launch._project_spec import LaunchProject
+
+from ..runner.abstract import AbstractRun
+from .run_queue_item_file_saver import RunQueueItemFileSaver
+
+
+@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