code: def memray_workers(
      directory: str | pathlib.Path = "memray-profiles",
      workers: int | None | list[str] = None,
      report_args: Sequence[str]
      | Literal[False] = ("flamegraph", "--temporal", "--leaks"),
      fetch_reports_parallel: bool | int = True,
      **memray_kwargs: Any,
  ) -> Iterator[None]:
      """Generate a Memray profile on the workers and download the generated report.
  
      Example::
  
          with memray_workers():
              client.submit(my_function).result()
  
          # Or even while the computation is already running
  
          fut = client.submit(my_function)
  
          with memray_workers():
              time.sleep(10)
  
          fut.result()
  
      Parameters
      ----------
      directory : str
          The directory to save the reports to.
      workers : int | None | list[str]
          The workers to profile. If int, the first n workers will be used.
          If None, all workers will be used.
          If list[str], the workers with the given addresses will be used.
      report_args : tuple[str]
          Particularly for native_traces=True, the reports have to be
          generated on the same host using the same Python interpreter as the
          profile was generated. Otherwise, native traces will yield unusable
          results. Therefore, we're generating the reports on the workers and
          download them afterwards. You can modify the report generation by
          providing additional arguments and we will generate the reports as::
  
              memray *report_args -f <filename> -o <filename>.html
  
          If the raw data should be fetched instead of the report, set this to
          False.
  
      fetch_reports_parallel : bool | int
          Fetching results is sometimes slow and it's sometimes not desired to
          wait for all workers to finish before receiving the first reports.
          This controls how many workers are fetched concurrently.
  
              int: Number of workers to fetch concurrently
              True: All workers concurrently
              False: One worker at a time
  
      **memray_kwargs
          Keyword arguments to be passed to memray.Tracker, e.g.
          {"native_traces": True}
      """
      directory = pathlib.Path(directory)
      client = get_client()
      scheduler_info = client.scheduler_info()
      worker_addr = scheduler_info["workers"]
      worker_names = {
          addr: winfo["name"]
          for addr, winfo in scheduler_info["workers"].items()
      }
      if not workers or isinstance(workers, int):
          nworkers = len(worker_addr)
          if isinstance(workers, int):
              nworkers = workers
          workers = list(worker_addr)[:nworkers]
      workers = list(workers)
      filename = uuid.uuid4().hex
      assert all(
          client.run(
              _start_memray, filename=filename, **memray_kwargs
          ).values()
      )
      # Sleep for a brief moment such that we get
      # a clear profiling signal when everything starts
      time.sleep(0.1)
      yield
      directory.mkdir(exist_ok=True)
  
      client = get_client()
      if fetch_reports_parallel is True:
          fetch_parallel = len(workers)
      elif fetch_reports_parallel is False:
          fetch_parallel = 1
      else:
          fetch_parallel = fetch_reports_parallel
  
      for w in partition(fetch_parallel, workers):
          try:
              profiles = client.run(
                  _fetch_memray_profile,
                  filename=filename,
                  report_args=report_args,
                  workers=w,
              )
              for worker_addr, profile in profiles.items():
                  path = directory / quote(
                      str(worker_names[worker_addr]), safe=""
                  )
                  if report_args:
                      suffix = ".html"
                  else:
                      suffix = ".memray"
                  with open(str(path) + suffix, "wb") as fd:
                      fd.write(profile)
  
          except Exception:
              logger.exception(
                  "Exception during report downloading from worker %s",
                  w,
              )
  
file_path: distributed/distributed/diagnostics/memray.py
imports: from __future__ import annotations
  import contextlib
  import logging
  import pathlib
  import subprocess
  import time
  import uuid
  from collections.abc import Iterator, Sequence
  from typing import Any, Literal
  from urllib.parse import quote
  from toolz.itertoolz import partition
  from distributed import get_client
  from distributed.worker import Worker
label_desc: Use image node_id 3 for calling a global function with example usage: memray_workers(directory, workers, report_args, fetch_reports_parallel) without return types