torchft-nightly 2026.1.3__cp310-cp310-manylinux_2_24_x86_64.whl

torchft/multiprocessing.py

@@ -0,0 +1,38 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import queue
+import time
+from datetime import timedelta
+from multiprocessing.connection import Connection
+from typing import Union
+
+import torch.multiprocessing as mp
+
+
+class _MonitoredPipe:
+    def __init__(self, pipe: "Connection[object, object]") -> None:  # type: ignore
+        self._pipe = pipe
+
+    def send(self, obj: object) -> None:
+        self._pipe.send(obj)
+
+    def recv(self, timeout: Union[float, timedelta]) -> object:
+        if isinstance(timeout, timedelta):
+            timeout = timeout.total_seconds()
+        if self._pipe.poll(timeout):
+            out = self._pipe.recv()
+            if isinstance(out, Exception):
+                raise out
+            return out
+        else:
+            raise TimeoutError(f"pipe.recv() timed out after {timeout} seconds")
+
+    def close(self) -> None:
+        self._pipe.close()
+
+    def closed(self) -> bool:
+        return self._pipe.closed

torchft/multiprocessing_dummy_context.py

@@ -0,0 +1,135 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Multiprocessing Dummy Context
+=========================
+
+This module provides a context-like interface for multiprocessing.dummy,
+which is a wrapper around the threading module that provides a multiprocessing-like
+interface but uses threads instead of processes.
+
+This allows code that uses multiprocessing.get_context() to work with
+multiprocessing.dummy by providing a compatible interface.
+"""
+
+import multiprocessing.dummy as mp
+import threading
+from typing import Callable, Iterable, Mapping
+
+
+class DummyContext:
+    """
+    A context-like class for multiprocessing.dummy that mimics the interface
+    of a context returned by multiprocessing.get_context().
+    """
+
+    def __init__(self, method: object = None) -> None:
+        """
+        Initialize the dummy context.
+
+        Args:
+            method: Ignored, only for compatibility with multiprocessing.get_context()
+        """
+        pass
+
+    def Process(
+        self,
+        group: object = None,
+        target: Callable[..., object] | None = None,
+        name: str | None = None,
+        args: Iterable[object] = (),
+        kwargs: Mapping[str, object] = {},
+        daemon: bool | None = None,
+    ) -> mp.DummyProcess:
+        """
+        Create a Process using multiprocessing.dummy.Process.
+        """
+        return mp.Process(
+            group=group, target=target, name=name, args=args, kwargs=kwargs
+        )
+
+    def Pipe(
+        self, duplex: bool = True
+    ) -> tuple[mp.connection.Connection, mp.connection.Connection]:
+        """
+        Create a Pipe using multiprocessing.dummy.Pipe.
+        """
+        return mp.Pipe(duplex)
+
+    def Queue(self, maxsize: int = 0) -> mp.Queue:
+        """
+        Create a Queue using multiprocessing.dummy.Queue.
+        """
+        return mp.Queue(maxsize)
+
+    def Event(self) -> threading.Event:
+        """
+        Create an Event using multiprocessing.dummy.Event.
+        """
+        return mp.Event()
+
+    def Lock(self) -> threading.Lock:
+        """
+        Create a Lock using multiprocessing.dummy.Lock.
+        """
+        return mp.Lock()
+
+    def RLock(self) -> threading.RLock:
+        """
+        Create an RLock using multiprocessing.dummy.RLock.
+        """
+        return mp.RLock()
+
+    def Semaphore(self, value: int = 1) -> threading.Semaphore:
+        """
+        Create a Semaphore using multiprocessing.dummy.Semaphore.
+        """
+        return mp.Semaphore(value)
+
+    def BoundedSemaphore(self, value: int = 1) -> threading.BoundedSemaphore:
+        """
+        Create a BoundedSemaphore using multiprocessing.dummy.BoundedSemaphore.
+        """
+        return mp.BoundedSemaphore(value)
+
+    def Condition(
+        self, lock: threading.Lock | threading.RLock | None = None
+    ) -> threading.Condition:
+        """
+        Create a Condition using multiprocessing.dummy.Condition.
+        """
+        return mp.Condition(lock)
+
+    def Manager(self) -> object:
+        """
+        Create a Manager using multiprocessing.dummy.Manager.
+        """
+        return mp.Manager()
+
+
+def get_context(method: object = None) -> DummyContext:
+    """
+    Return a context object for multiprocessing.dummy.
+
+    This function mimics multiprocessing.get_context() but returns a DummyContext
+    that works with multiprocessing.dummy. This can be used to patch
+    multiprocessing.dummy like so
+
+
+    ```
+    import multiprocessing.dummy as mp
+    from torchft.multiprocessing_dummy_context import get_context
+    mp.get_context = get_context
+    ```
+
+    Args:
+        method: Ignored, only for compatibility with multiprocessing.get_context()
+
+    Returns:
+        A DummyContext instance
+    """
+    return DummyContext(method)

torchft/multiprocessing_test.py

@@ -0,0 +1,58 @@
+from multiprocessing.connection import Connection
+from unittest import TestCase
+
+import torch.multiprocessing as mp
+
+from torchft.multiprocessing import _MonitoredPipe
+
+
+def pipe_get(q: "Connection[object, object]") -> None:  # type: ignore
+    q.recv()
+
+
+def pipe_put(q: "Connection[object, object]") -> None:  # type: ignore
+    q.recv()
+    q.send(1)
+
+
+class MultiprocessingTest(TestCase):
+    def test_monitored_queue_put(self) -> None:
+        ctx = mp.get_context("fork")
+        local, remote = ctx.Pipe()
+        p = ctx.Process(target=pipe_get, args=(remote,), daemon=True)
+        p.start()
+        del remote
+
+        mq = _MonitoredPipe(local)
+        mq.send(1)
+        with self.assertRaisesRegex(
+            (ConnectionResetError, BrokenPipeError),
+            "(Connection reset by peer|Broken pipe)",
+        ):
+            while True:
+                mq.send(1)
+
+        mq.close()
+        assert mq.closed()
+
+    def test_monitored_queue_get(self) -> None:
+        ctx = mp.get_context("fork")
+        local, remote = ctx.Pipe()
+        p = ctx.Process(target=pipe_put, args=(remote,), daemon=True)
+        p.start()
+        del remote
+
+        mq = _MonitoredPipe(local)
+
+        with self.assertRaisesRegex(TimeoutError, "timed out after 0.0 seconds"):
+            mq.recv(timeout=0.0)
+
+        # continue
+        mq.send(1)
+
+        self.assertEqual(mq.recv(timeout=10), 1)
+        with self.assertRaises(EOFError):
+            mq.recv(timeout=10)
+
+        mq.close()
+        assert mq.closed()

torchft/optim.py

@@ -0,0 +1,63 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Optimizers
+============
+
+This module implements an optimizer wrapper that works with the Manager to provide fault tolerance.
+
+"""
+
+from typing import Any, Dict, List, Mapping, Optional, TYPE_CHECKING
+
+import torch
+from torch.optim import Optimizer
+
+if TYPE_CHECKING:
+    from torchft.manager import Manager
+
+
+class OptimizerWrapper(Optimizer):
+    """
+    This wraps any provided optimizer and in conjunction with the manager will provide fault tolerance.
+
+    zero_grad() must be called at the start of the forwards pass and step() must
+    be called at the end of the backwards pass.
+
+    Depending on the state of the manager, the optimizer will either commit the
+    gradients to the wrapped optimizer or ignore them.
+    """
+
+    def __init__(self, manager: "Manager", optim: Optimizer) -> None:
+        self.optim = optim
+        self.manager = manager
+
+    def add_param_group(self, param_group: Dict[str, Any]) -> None:
+        self.optim.add_param_group(param_group)
+
+    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
+        self.optim.load_state_dict(state_dict)
+
+    def state_dict(self) -> Dict[str, Any]:
+        return self.optim.state_dict()
+
+    def zero_grad(self, set_to_none: bool = True) -> None:
+        self.manager.start_quorum()
+        self.optim.zero_grad(set_to_none)
+
+    def step(self, closure: Optional[object] = None) -> None:
+        assert closure is None, "optimizers that use closures are not supported"
+        if self.manager.should_commit():
+            self.optim.step()
+
+    @property
+    def param_groups(self) -> List[Dict[str, Any]]:
+        return self.optim.param_groups
+
+    @property
+    def state(self) -> Mapping[torch.Tensor, object]:
+        return self.optim.state

torchft/optim_test.py

@@ -0,0 +1,50 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from unittest import TestCase
+from unittest.mock import create_autospec, MagicMock
+
+import torch
+from torch.nn import Linear
+from torch.optim import AdamW
+
+from torchft.manager import Manager
+from torchft.optim import OptimizerWrapper
+
+
+class TestOptim(TestCase):
+    def test_optimizer_wrapper(self) -> None:
+        manager = create_autospec(Manager)
+
+        m = Linear(3, 4)
+        base_optim = AdamW(m.parameters())
+        optim = OptimizerWrapper(manager, base_optim)
+        optim.add_param_group(
+            {
+                "params": [],
+                "lr": 1e-4,
+            }
+        )
+
+        # test state_dict handling
+        optim.load_state_dict(optim.state_dict())
+
+        optim.zero_grad()
+        self.assertEqual(manager.start_quorum.call_count, 1)
+
+        b = torch.rand(3)
+        m(b).sum().backward()
+
+        manager.should_commit.return_value = True
+        optim.step()
+        manager.should_commit.return_value = False
+        optim.step()
+        self.assertEqual(len(optim.param_groups), 2)
+        self.assertEqual(optim.param_groups[1]["lr"], 1e-4)
+        self.assertEqual(optim.param_groups[1]["params"], [])
+        self.assertEqual(len(optim.state), len(list(m.parameters())))
+
+        self.assertEqual(manager.should_commit.call_count, 2)

torchft/otel.py

@@ -0,0 +1,134 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import json
+import logging
+import os
+import time
+from typing import Any, List, Sequence, TYPE_CHECKING
+
+from opentelemetry._logs import set_logger_provider
+
+from opentelemetry.exporter.otlp.proto.http._log_exporter import OTLPLogExporter
+from opentelemetry.sdk._logs import LoggerProvider, LoggingHandler
+from opentelemetry.sdk._logs.export import BatchLogRecordProcessor
+from opentelemetry.sdk.resources import Resource
+
+# These types are available in opentelemetry-sdk but Pyre's type stubs
+# don't include them. We import them at runtime and provide type aliases for
+# static type checking.
+if TYPE_CHECKING:
+    # pyre-fixme[33]: Aliasing to Any is prohibited. opentelemetry-sdk lacks type stubs.
+    ReadableLogRecord = Any
+    # pyre-fixme[33]: Aliasing to Any is prohibited. opentelemetry-sdk lacks type stubs.
+    LogRecordExporter = Any
+    # pyre-fixme[33]: Aliasing to Any is prohibited. opentelemetry-sdk lacks type stubs.
+    LogRecordExportResult = Any
+    # pyre-fixme[33]: Aliasing to Any is prohibited. opentelemetry-sdk lacks type stubs.
+    ConsoleLogRecordExporter = Any
+else:
+    from opentelemetry.sdk._logs import ReadableLogRecord
+    from opentelemetry.sdk._logs.export import (
+        ConsoleLogRecordExporter,
+        LogRecordExporter,
+        LogRecordExportResult,
+    )
+
+_LOGGER_PROVIDER: dict[str, LoggerProvider] = {}
+# Path to the file containing OTEL resource attributes
+TORCHFT_OTEL_RESOURCE_ATTRIBUTES_JSON = "TORCHFT_OTEL_RESOURCE_ATTRIBUTES_JSON"
+
+
+class TeeLogExporter(LogRecordExporter):
+    """Exporter that writes to multiple exporters."""
+
+    def __init__(
+        self,
+        exporters: List[LogRecordExporter],
+    ) -> None:
+        self._exporters = exporters
+
+    def export(self, batch: Sequence[ReadableLogRecord]) -> LogRecordExportResult:
+        for e in self._exporters:
+            e.export(batch)
+        return LogRecordExportResult.SUCCESS
+
+    def shutdown(self) -> None:
+        for e in self._exporters:
+            e.shutdown()
+
+
+def setup_logger(name: str) -> None:
+    if os.environ.get("TORCHFT_USE_OTEL", "false") == "false":
+        return
+
+    if name in _LOGGER_PROVIDER:
+        return
+
+    torchft_otel_resource_attributes_json = os.environ.get(
+        TORCHFT_OTEL_RESOURCE_ATTRIBUTES_JSON
+    )
+
+    if torchft_otel_resource_attributes_json is not None:
+        with open(torchft_otel_resource_attributes_json) as f:
+            attributes = json.loads(f.read())
+            resource = Resource.create(attributes=attributes[name])
+    else:
+        resource = Resource.create()
+
+    logger_provider = LoggerProvider(resource=resource)
+    set_logger_provider(logger_provider)
+
+    exporter = TeeLogExporter(
+        exporters=[
+            ConsoleLogRecordExporter(),
+            OTLPLogExporter(
+                timeout=5,
+            ),
+        ],
+    )
+    logger_provider.add_log_record_processor(BatchLogRecordProcessor(exporter))
+    handler = LoggingHandler(level=logging.NOTSET, logger_provider=logger_provider)
+
+    # Attach OTLP handler to otel logger
+    logging.getLogger(name).addHandler(handler)
+
+    _LOGGER_PROVIDER[name] = logger_provider
+
+
+def shutdown() -> None:
+    for logger_provider in _LOGGER_PROVIDER.values():
+        logger_provider.shutdown()
+
+
+# Example usage of the logger
+def main() -> None:
+    logging.getLogger().setLevel(logging.INFO)
+    setup_logger("torchft_test")
+
+    while True:
+        time.sleep(1)
+        loggers = [
+            logging.getLogger("torchft_test"),
+            logging.getLogger("myapp.area1"),
+            logging.getLogger(),
+        ]
+
+        for i, logger in enumerate(loggers):
+            # only this should be picked up by OTEL when using otel logger
+            logger.info(
+                "Quick zephyrs blow, vexing daft Jim.",
+                extra={
+                    "test_attr": f"value{i}",
+                },
+            )
+            logger.debug("Jackdaws love my big sphinx of quartz.")
+
+        print("Example done; exiting...")
+
+
+if __name__ == "__main__":
+    main()

torchft/parameter_server.py

@@ -0,0 +1,195 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Parameter Servers
+==================
+
+This module provides a prototype implementation of a fault tolerant parameter server bulit on the reconfigurable ProcessGroups.
+"""
+
+import json
+import logging
+import socket
+import threading
+import urllib.request
+import uuid
+from abc import ABC, abstractmethod
+from http.server import BaseHTTPRequestHandler
+
+from torch.distributed import TCPStore
+
+from torchft.http import _IPv6HTTPServer
+from torchft.process_group import ProcessGroup
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+
+class ParameterServer(ABC):
+    """
+    This implements a threaded parameter server using the torchft reconfigurable
+    ProcessGroups.
+    """
+
+    def __init__(self, port: int, store_port: int = 0) -> None:
+        """
+        Create a new ParameterServer.
+
+        Args:
+            port: the port to bind the HTTP server to.
+            store_port: the port to bind the TCPStore server to.
+        """
+        self.store = TCPStore(
+            host_name="0.0.0.0",
+            port=store_port,
+            is_master=True,
+            wait_for_workers=False,
+        )
+
+        ps = self
+
+        class RequestHandler(BaseHTTPRequestHandler):
+            def do_GET(self):
+                if self.path != "/new_session":
+                    self.send_response(400)
+                    self.send_header("Content-type", "text/plain")
+                    self.end_headers()
+                    self.err(f"invalid path, got {self.path}")
+                    return
+
+                try:
+                    self.send_response(200)
+                    self.send_header(
+                        "Content-type", "application/json"
+                    )  # TODO: correct mime type
+                    self.end_headers()
+
+                    session_id = str(uuid.uuid4())
+
+                    store_addr = (
+                        f"{socket.gethostname()}:{ps.store.port}/session/{session_id}"
+                    )
+
+                    logger.info(f"creating new session {session_id}")
+
+                    data = (
+                        json.dumps(
+                            {
+                                "session_id": session_id,
+                                "store_addr": store_addr,
+                            }
+                        )
+                        + "\n"
+                    )
+                    data = data.encode()
+
+                    self.wfile.write(data)
+
+                    # close the connection up front so client will know json is
+                    # complete
+                    self.finish()
+                    self.connection.close()
+
+                    # hijack thread for the session
+                    ps._handle_session(session_id, store_addr)
+                except Exception:
+                    logger.exception(
+                        f"got exception in request handler for {self.path}"
+                    )
+                    raise
+
+        server_address = ("", port)
+        self._server = _IPv6HTTPServer(server_address, RequestHandler)
+        self._server.daemon_threads = True
+        logger.info(f"Started ParameterServer on {self.address()}...")
+
+        self._thread = threading.Thread(
+            target=self._serve,
+            args=(),
+            daemon=True,
+        )
+        self._thread.start()
+
+    def address(self) -> str:
+        """
+        Returns the HTTP address to create a new session on this server.
+
+        Format: http://host:port/new_session
+
+        Returns:
+            an HTTP address
+        """
+        port = self._server.socket.getsockname()[1]
+        return f"http://{socket.gethostname()}:{port}/new_session"
+
+    def _serve(self) -> None:
+        try:
+            self._server.serve_forever()
+        except Exception as e:
+            logger.exception("got exception in checkpoint server")
+
+    @classmethod
+    @abstractmethod
+    def new_process_group(cls) -> ProcessGroup:
+        """
+        Create a new non-configured ProcessGroup for the ParameterServer to
+        configure when setting up server and client connections.
+
+        Must be implemented by subclasses.
+
+        Returns:
+            a new ProcessGroup
+        """
+        ...
+
+    @classmethod
+    def new_session(cls, address: str) -> ProcessGroup:
+        """
+        Creates a new session on the parameter server and returns a ProcessGroup
+        configured for that server.
+
+        Client is rank 1, server is rank 0.
+        """
+        with urllib.request.urlopen(address) as f:
+            data = json.load(f)
+
+        session_id = data["session_id"]
+        store_addr = data["store_addr"]
+
+        logger.info(f"connecting to session {session_id} at {store_addr}")
+
+        pg = cls.new_process_group()
+        # client is always rank 1
+        pg.configure(store_addr, replica_id="0", rank=1, world_size=2)
+
+        return pg
+
+    def _handle_session(self, session_id: str, store_addr: str) -> None:
+        pg = self.new_process_group()
+        # paramter server is always rank 0
+        pg.configure(store_addr, replica_id="0", rank=0, world_size=2)
+
+        self.forward(session_id, pg)
+
+    @abstractmethod
+    def forward(self, session_id: str, pg: ProcessGroup) -> None:
+        """
+        This method will be called once per session in a dedicated thread. To
+        support multiple operations on a single session you should put a
+        for-loop in your forward implementation.
+
+        If an error occurs, the process group will be freed and the client will
+        have to create a new session.
+
+        The server rank is 0 and the client rank is 1.
+
+        Must be implemented by subclasses.
+
+        Args:
+            session_id: a unique uuid for this session
+            pg: the ProcessGroup that's configured for the client.
+        """
+        ...