agent 2 sweep ini outline

wandb/sdk/launch/agent2/agent.py

@@ -11,7 +11,10 @@
 from wandb.sdk.launch.agent.agent import HIDDEN_AGENT_RUN_TYPE
 from wandb.sdk.launch.agent.job_status_tracker import JobAndRunStatusTracker
 from wandb.sdk.launch.agent.run_queue_item_file_saver import RunQueueItemFileSaver
+from wandb.sdk.launch.agent2.controllers.local_process import SchedulerManager
 from wandb.sdk.launch.builder.build import construct_agent_configs
+from wandb.sdk.launch.environment.local_environment import LocalEnvironment
+from wandb.sdk.launch.registry.local_registry import LocalRegistry
 from wandb.sdk.launch.utils import PROJECT_SYNCHRONOUS, event_loop_thread_exec
 
 from .controller import LaunchController, LegacyResources
@@ -68,6 +71,7 @@
  self._last_state = None
  self._wandb_version: str = "wandb@" + wandb.__version__
  self._task: Optional[asyncio.Task[Any]] = None
+ self._receive_scheduler_job_queue = asyncio.Queue()
 
  self._logger = logging.getLogger("wandb.launch.agent2")
  handler = logging.StreamHandler(sys.stdout)
@@ -123,7 +127,67 @@
  # Start the main agent state poll loop
  self.start_poll_loop(event_loop)
 
+ def file_saver_factory(job_id):
+ return RunQueueItemFileSaver(self._wandb_run, job_id)
+
+ def job_tracker_factory(job_id, q):
+ return JobAndRunStatusTracker(job_id, q, file_saver_factory(job_id))
+
  try:
+ # create sweep scheduler local process controller
+ # TODO: move into util function
+ controller_impl = self.get_controller_for_jobset("local-process")
+ _, build_config, registry_config = construct_agent_configs(
+ dict(self._config)
+ )
+ environment = LocalEnvironment()
+ registry = LocalRegistry()
+ runner = loader.runner_from_config(
+ "local-process",
+ self._api, # todo factor out (?)
+ {},
+ environment,
+ registry,
+ )
+ legacy_resources = LegacyResources(
+ self._api,
+ builder,
+ registry,
+ runner,
+ environment,
+ job_tracker_factory,
+ )
+ controller_logger = self._logger.getChild(
+ "controller.sweep-scheduler-local-process"
+ )
+ scheduler_controller = controller_impl(
+ {
+ "agent_id": self._id,
+ "jobset_spec": JobSetSpec(
+ name="_wandb_sweep_scheduler_local_process",
+ entity_name=self._config["entity"],
+ project_name="_wandb_sweep-scheduler_local_process",
+ ),
+ "jobset_metadata": None,
+ },
+ JobSet(self._api, {}, self._id, controller_logger),
+ controller_logger,
+ self._shutdown_controllers_event,
+ legacy_resources,
+ self._receive_scheduler_job_queue, # TODO: not necessary for sweep scheduler
+ )
+ manager_logger = self._logger.getChild("scheduler_manager")
+ scheduler_manager = SchedulerManager(
+ scheduler_controller,
+ self._config["max_schedulers"],
+ self._receive_scheduler_job_queue,
+ manager_logger,
+ )
+ controller_task: asyncio.Task = asyncio.create_task(
+
+ )
+ self._launch_controller_tasks.add(controller_task)
+
  # Start job set and controller loops
  for q in self._config["queues"]:
  # Start a JobSet for each queue
@@ -169,12 +233,6 @@
  registry,
  )
 
- def file_saver_factory(job_id):
- return RunQueueItemFileSaver(self._wandb_run, job_id)
-
- def job_tracker_factory(job_id, q=q):
- return JobAndRunStatusTracker(job_id, q, file_saver_factory(job_id))
-
  legacy_resources = LegacyResources(
  self._api,
  builder,
@@ -197,6 +255,7 @@
  controller_logger,
  self._shutdown_controllers_event,
  legacy_resources,
+ self._receive_scheduler_job_queue,
  )
  )
  self._launch_controller_tasks.add(controller_task)

wandb/sdk/launch/agent2/controllers/base.py

@@ -14,11 +14,11 @@
 from wandb.sdk.launch.errors import LaunchError
 from wandb.sdk.launch.runner.abstract import AbstractRun, Status
 
-from ...agent.agent import RUN_INFO_GRACE_PERIOD
+from ...agent.agent import RUN_INFO_GRACE_PERIOD, _is_scheduler_job
 from ...queue_driver.abstract import AbstractQueueDriver
 from ...utils import event_loop_thread_exec
 from ..controller import LaunchControllerConfig, LegacyResources
-from ..jobset import Job, JobSet
+from ..jobset import Job, JobSet, JobWithQueue
 
 WANDB_JOBSET_DISCOVERABILITY_LABEL = "_wandb-jobset"
 
@@ -41,13 +41,15 @@
  jobset: JobSet,
  logger: logging.Logger,
  legacy: LegacyResources,
+ scheduler_queue: asyncio.Queue[Tuple[JobWithQueue, asyncio.Future]],
  max_concurrency: int,
  ):
  self.config = config
  self.jobset = jobset
  self.logger = logger
  self.legacy = legacy
  self.max_concurrency = max_concurrency
+ self._scheduler_queue = scheduler_queue
 
  self.id = config["jobset_spec"].name
  self.active_runs: Dict[str, RunWithTracker] = {}
@@ -126,7 +128,7 @@
  try:
  project = LaunchProject.from_spec(job.run_spec, self.legacy.api)
  run_id = project.run_id
- job_tracker = self.legacy.job_tracker_factory(run_id)
+ job_tracker = self.legacy.job_tracker_factory(run_id, project.queue_name)
  job_tracker.update_run_info(project)
  except Exception as e:
  self.logger.error(
@@ -140,6 +142,16 @@
  project.run_queue_item_id = job.id
  project.fetch_and_validate_project()
 
+ if (
+ _is_scheduler_job(job.run_spec)
+ and job.run_spec.get("resource") == "local-process"
+ ):
+ future = asyncio.futures.Future()
+ await self._scheduler_queue.put((job, future))
+ res = await future.result()
+ if res == False:
+ return None
+
  ack_result = await self.jobset.ack_job(job.id, run_id)
  self.logger.info(f"Acked item: {json.dumps(ack_result, indent=2)}")
  if not ack_result:

wandb/sdk/launch/agent2/controllers/local_process.py

@@ -1,12 +1,12 @@
 import asyncio
 import json
 import logging
-from typing import Any, List, Optional
+from typing import Any, List, Optional, Tuple, Union
 
 from ..._project_spec import LaunchProject
 from ...queue_driver import passthrough
 from ..controller import LaunchControllerConfig, LegacyResources
-from ..jobset import Job, JobSet
+from ..jobset import Job, JobSet, JobWithQueue
 from .base import BaseManager, RunWithTracker
 
 
@@ -16,6 +16,7 @@
  logger: logging.Logger,
  shutdown_event: asyncio.Event,
  legacy: LegacyResources,
+ agent_queue: asyncio.Queue,
 ) -> Any:
  # disable job set loop because we are going to use the passthrough queue driver
  # to drive the launch controller here
@@ -39,7 +40,9 @@
  f"Starting local process controller with max concurrency {max_concurrency}"
  )
 
- mgr = LocalProcessManager(config, jobset, logger, legacy, max_concurrency)
+ mgr = LocalProcessManager(
+ config, jobset, logger, legacy, agent_queue, max_concurrency
+ )
 
  while not shutdown_event.is_set():
  await mgr.reconcile()
@@ -62,6 +65,7 @@
  jobset: JobSet,
  logger: logging.Logger,
  legacy: LegacyResources,
+ agent_queue: asyncio.Queue,
  max_concurrency: int,
  ):
  self.queue_driver: passthrough.PassthroughQueueDriver = (
@@ -73,7 +77,7 @@
  agent_id=config["agent_id"],
  )
  )
- super().__init__(config, jobset, logger, legacy, max_concurrency)
+ super().__init__(config, jobset, logger, legacy, agent_queue, max_concurrency)
 
  async def reconcile(self) -> None:
  num_runs_needed = self.max_concurrency - len(self.active_runs)
@@ -90,35 +94,99 @@
  async def launch_item(self, item: Job) -> Optional[str]:
  self.logger.info(f"Launching item: {item}")
 
- project = LaunchProject.from_spec(item.run_spec, self.legacy.api)
- project.queue_name = self.config["jobset_spec"].name
- project.queue_entity = self.config["jobset_spec"].entity_name
- project.run_queue_item_id = item.id
+ project = self._populate_project(item)
  project.fetch_and_validate_project()
+ run_id = await self._launch_job(item, project)
+ self.logger.info(f"Launched item got run_id: {run_id}")
+ return run_id
+
+ async def launch_scheduler_item(self, item: JobWithQueue) -> Optional[str]:
+ self.logger.info(f"Launching item: {item}")
 
+ project = self._populate_project(item)
+ project.fetch_and_validate_project()
+
+ run_id = await self._launch_job(item.job, project)
+ self.logger.info(f"Launched item got run_id: {run_id}")
+ return run_id
+
+ def _populate_project(self, job: Union[Job, JobWithQueue]) -> LaunchProject:
+ project = None
+ if isinstance(job, JobWithQueue):
+ project = LaunchProject.from_spec(job.job.run_spec, self.legacy.api)
+ queue_name = job.queue
+ queue_entity = job.entity
+ job_id = job.job.id
+ else:
+ project = LaunchProject.from_spec(job.run_spec, self.legacy.api)
+ queue_name = self.config["jobset_spec"].name
+ queue_entity = self.config["jobset_spec"].entity_name
+ job_id = job.id
+ project.queue_name = queue_name
+ project.queue_entity = queue_entity
+ project.run_queue_item_id = job_id
+ return project
+
+ def _get_job(self, item: Union[Job, JobWithQueue]) -> Job:
+ if isinstance(item, JobWithQueue):
+ return item.job
+ return item
+
+ async def _launch_job(self, job: Job, project: LaunchProject) -> Optional[str]:
  run_id = project.run_id
- job_tracker = self.legacy.job_tracker_factory(run_id)
+ job_tracker = self.legacy.job_tracker_factory(run_id, project.queue_name)
  job_tracker.update_run_info(project)
 
- ack_result = await self.queue_driver.ack_run_queue_item(item.id, run_id)
+ # note since we ack on rqi id the queue driver will handle acking the run queue item
+ # even if its not for the specified queue
+ ack_result = await self.queue_driver.ack_run_queue_item(job.id, run_id)
  if ack_result is None:
- self.logger.error(f"Failed to ack item {item.id}")
+ self.logger.error(f"Failed to ack item {job.id}")
  return None
  self.logger.info(f"Acked item: {json.dumps(ack_result, indent=2)}")
  run = await self.legacy.runner.run(project, "") # image is unused
  if not run:
  job_tracker.failed_to_start = True
- self.logger.error(f"Failed to start run for item {item.id}")
+ self.logger.error(f"Failed to start run for item {job.id}")
  raise NotImplementedError("TODO: handle this case")
 
- self.active_runs[item.id] = RunWithTracker(run, job_tracker)
+ self.active_runs[job.id] = RunWithTracker(run, job_tracker)
 
  run_id = project.run_id
- self.logger.info(f"Launched item got run_id: {run_id}")
  return run_id
 
  async def find_orphaned_jobs(self) -> List[Any]:
  raise NotImplementedError
 
  def label_job(self, project: LaunchProject) -> None:
  pass
+
+
+class SchedulerManager:
+ def __init__(
+ self,
+ controller: LocalProcessManager,
+ max_jobs: int,
+ scheduler_jobs_queue: asyncio.Queue[Tuple[JobWithQueue, asyncio.Future]],
+ logger: logging.Logger,
+ ):
+ self._controller = controller
+ self._scheduler_jobs_queue = scheduler_jobs_queue
+ self._logger = logger
+ self._max_jobs = max_jobs
+
+ async def poll(self):
+ while True:
+ res = await self._scheduler_jobs_queue.get()
+ if res is None:
+ asyncio.sleep(5) # TODO: const this
+ break
+ job, future = res
+ if len(self._controller.active_runs) >= self._max_jobs:
+ self._logger.info(f"Scheduler job queue is full, skipping job: {job}")
+ future.set_result(False)
+ continue
+ future.set_result(True)
+ asyncio.create_task(self.controller.launch_scheduler_item(job))
+ self._scheduler_jobs_queue.task_done()
+ self._logger.info(f"Launched scheduler job: {job}")

wandb/sdk/launch/agent2/jobset.py

@@ -48,6 +48,13 @@ class Job:
  claimed_by: str
 
 
+@dataclass
+class JobWithQueue:
+ job: Job
+ queue: str
+ entity: str
+
+
 def run_queue_item_to_job(run_queue_item: Dict[str, Any]) -> Job:
  return Job(
  id=run_queue_item["id"],