Merge branch 'apple:main' into jax_0.9.1_dev

Author: andersensam

axlearn/audio/frontend_utils.py

@@ -15,13 +15,11 @@
 import jax
 import jax.numpy as jnp
 import numpy as np
-from jax._src.mesh import thread_resources
-from jax.experimental.shard_map import shard_map
 from jax.sharding import PartitionSpec
 from numpy.typing import ArrayLike
 
 from axlearn.common import ein_ops
-from axlearn.common.utils import Tensor
+from axlearn.common.utils import Tensor, get_current_abstract_or_physical_mesh
 
 
 class WindowType(enum.Enum):
@@ -402,9 +400,9 @@ def sharded_fft(n: int, partition_spec: PartitionSpec) -> Callable[[Tensor], Ten
     Returns:
         A callable that computes FFT.
     """
-    return shard_map(
+    return jax.shard_map(
         lambda x: jnp.fft.rfft(cast_for_rfft(x), n=n),
-        mesh=thread_resources.env.physical_mesh,
+        mesh=get_current_abstract_or_physical_mesh(),
         in_specs=partition_spec,
         out_specs=partition_spec,
     )

axlearn/cloud/common/bastion.py

@@ -110,6 +110,7 @@
 except ModuleNotFoundError:
     logging.warning("tensorflow_io is not installed -- tf_io may not work with s3://")
 
+from axlearn.cloud.common import metrics
 from axlearn.cloud.common.cleaner import Cleaner
 from axlearn.cloud.common.event_queue import BaseQueueClient, Event
 from axlearn.cloud.common.quota import QuotaFn
@@ -904,10 +905,13 @@ def check_leaves(x, y):
 
     def _wait_and_close_proc(self, proc: _PipedProcess, kill: bool = False):
         """Cleans up the process/fds and upload logs to gs."""
+        pid = proc.popen.pid
         if kill:
             send_signal(proc.popen, sig=signal.SIGKILL)
         # Note: proc should already be polled and completed, so wait is nonblocking.
         proc.popen.wait()
+        # Cleanup the local prometheus DB files when the process is terminated
+        _catch_with_error_log(metrics.cleanup, pid)
         proc.fd.close()
         # Upload outputs to log dir.
         _catch_with_error_log(
@@ -958,15 +962,16 @@ def _sync_jobs(self):
 
         We use these jobspecs to update the local self._active_jobs.
         """
-        active_jobs, jobs_with_user_states, jobs_with_failed_validation = download_job_batch(
-            spec_dir=self._active_dir,
-            state_dir=self._state_dir,
-            user_state_dir=self._user_state_dir,
-            verbose=True,
-            remove_invalid_user_states=True,
-            quota=self._quota,
-            validator=self._validator,
-        )
+        with metrics.BASTION_DOWNLOAD_JOB_ALL_SECONDS.time():
+            active_jobs, jobs_with_user_states, jobs_with_failed_validation = download_job_batch(
+                spec_dir=self._active_dir,
+                state_dir=self._state_dir,
+                user_state_dir=self._user_state_dir,
+                verbose=True,
+                remove_invalid_user_states=True,
+                quota=self._quota,
+                validator=self._validator,
+            )
         self._jobs_with_user_states = jobs_with_user_states
         # Iterate over unique job names.
         # pylint: disable-next=use-sequence-for-iteration

axlearn/cloud/common/metrics.py

@@ -0,0 +1,158 @@
+"""Prometheus metrics for Bastion runners."""
+
+import logging
+import threading
+from http.server import BaseHTTPRequestHandler, HTTPServer
+
+from prometheus_client import CollectorRegistry, Gauge, Histogram, generate_latest, multiprocess
+
+BASTION_DOWNLOAD_JOB_ALL_SECONDS = Histogram(
+    "bastion_download_job_all_seconds",
+    "Time in seconds consumed in each scheduling loop to download all the job specs",
+    [],
+    buckets=(1, 5, 10, 30, 60, 120, 300, 600, 1800, 3600),  # 1s to 1hr
+)
+
+JOB_TIME_TO_RUNNING_SECONDS = Histogram(
+    "bastion_job_time_to_running_seconds",
+    "Time in seconds from runner starts a job until until it reaches Running state",
+    ["job_name"],
+    buckets=(1, 5, 10, 30, 60, 120, 300, 600, 1800, 3600),  # 1s to 1hr
+)
+
+JOB_RUN_LATENCY_GAUGE = Gauge(
+    "bastion_job_run_latency_seconds",
+    "Time in seconds from runner starts a job until until it reaches Running state",
+    ["job_name"],
+)
+
+JOB_WAIT_FOR_BUILD_SECONDS = Gauge(
+    "bastion_job_wait_for_build_seconds",
+    "time in seconds until the job is built successfully",
+    ["job_name"],
+)
+
+JOB_SUSPENDED_DURATION_SECONDS = Histogram(
+    "bastion_job_suspended_duration_seconds",
+    "Time in seconds the jobset stayed in SUSPENDED state",
+    ["job_name"],
+    buckets=(1, 5, 10, 30, 60, 120, 300, 600, 1800, 3600),  # 1s to 1hr
+)
+
+
+class _MetricsHandler(BaseHTTPRequestHandler):
+    """HTTP handler that serves aggregated metrics from multiprocess mode."""
+
+    def do_GET(self):  # pylint: disable=invalid-name
+        if self.path == "/metrics":
+            registry = CollectorRegistry()
+            multiprocess.MultiProcessCollector(registry)
+            metrics = generate_latest(registry)
+
+            self.send_response(200)
+            self.send_header("Content-Type", "text/plain; charset=utf-8")
+            self.end_headers()
+            self.wfile.write(metrics)
+        else:
+            self.send_error(404)
+
+    def log_message(self, format, *args):  # pylint: disable=redefined-builtin
+        # Suppress default logging
+        pass
+
+
+class MetricsServer:
+    """Singleton class to manage Prometheus metrics server.
+
+    Supports multiprocess mode where metrics from parent and child processes
+    are aggregated and exposed via a single endpoint.
+    """
+
+    _instance = None
+    _lock = threading.Lock()
+    _server_started = False
+    _multiprocess_dir = None
+
+    def __new__(cls):
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    cls._instance = super().__new__(cls)
+        return cls._instance
+
+    @classmethod
+    def start(cls, port: int = 8000):
+        """Start Prometheus metrics HTTP server on the specified port.
+
+        Args:
+            port: Port number to start the metrics server on. Defaults to 8000.
+            use_multiprocess: If True, use multiprocess mode which aggregates metrics
+                            from child processes. Defaults to False for backward compatibility.
+        """
+        with cls._lock:
+            if not cls._server_started:
+                try:
+                    # Use custom handler for multiprocess aggregation
+                    server = HTTPServer(("", port), _MetricsHandler)
+                    thread = threading.Thread(target=server.serve_forever, daemon=True)
+                    thread.start()
+                    logging.info(
+                        "Prometheus metrics server (multiprocess mode) started on port %s", port
+                    )
+                    cls._server_started = True
+                except OSError as e:
+                    # Port might already be in use
+                    logging.warning("Failed to start metrics server on port %s: %s", port, e)
+
+    @classmethod
+    def is_started(cls) -> bool:
+        """Check if the metrics server has been started."""
+        return cls._server_started
+
+    @classmethod
+    def get_multiprocess_dir(cls) -> str:
+        """Get the multiprocess directory path if configured."""
+        return cls._multiprocess_dir
+
+
+def record_job_time_to_running(job_name: str, time_seconds: float):
+    """Record the time it took for a job to reach RUNNING state.
+
+    Args:
+        job_name: Name of the job
+        time_seconds: Time in seconds from first seen until RUNNING state
+    """
+    JOB_TIME_TO_RUNNING_SECONDS.labels(job_name=job_name).observe(time_seconds)
+    logging.info("Metrics: job %s took %.2f seconds to reach RUNNING state", job_name, time_seconds)
+
+
+def record_job_run_latency(job_name: str, time_seconds: float):
+    JOB_RUN_LATENCY_GAUGE.labels(job_name).set(time_seconds)
+
+
+def record_job_wait_for_build(job_name: str, time_seconds: float):
+    """Record the time it took to wait for the cloud build to finish
+
+    This does not necessarily mean the whole cloud build time for a job,
+    as the build starts when the client submits the job. This only measures
+    how long the runner waits for the build to complete before launching
+    the job.
+    """
+    JOB_WAIT_FOR_BUILD_SECONDS.labels(job_name).set(time_seconds)
+
+
+def record_job_suspended_duration(job_name: str, time_seconds: float):
+    """Record the time a job spent in SUSPENDED state.
+
+    Args:
+        job_name: Name of the job
+        time_seconds: Time in seconds the job was in SUSPENDED state
+    """
+    JOB_SUSPENDED_DURATION_SECONDS.labels(job_name=job_name).observe(time_seconds)
+    logging.info("Metrics: job %s was in SUSPENDED state for %.2f seconds", job_name, time_seconds)
+
+
+def cleanup(pid=None):
+    if pid:
+        logging.info("Cleaning up prometheus metrics DB file for %s", pid)
+        multiprocess.mark_process_dead(pid)