modelexpress-client 0.6.0

Client library for Model Express gRPC server
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

from unittest.mock import MagicMock, patch

import modelexpress_rl
from modelexpress_rl.train.resources import _TrainerResources


def test_resources_initialize_transport_and_manifest_service(monkeypatch):
    monkeypatch.setenv("MX_WORKER_HOST", "trainer.test")
    monkeypatch.setenv("MX_METADATA_PORT", "18000")
    monkeypatch.setenv("MX_WORKER_GRPC_PORT", "19000")
    monkeypatch.setenv("RANK", "7")
    manager = MagicMock()
    server = MagicMock()
    server.add_insecure_port.return_value = 19002

    with (
        patch(
            "modelexpress.nixl_transfer.NixlTransferManager",
            return_value=manager,
        ) as manager_type,
        patch(
            "modelexpress_rl.train.resources.grpc.server",
            return_value=server,
        ),
    ):
        resources = _TrainerResources.initialize(device_id=2)

    manager_type.assert_called_once_with(
        agent_name="modelexpress-trainer-7",
        device_id=2,
        listen_port=18002,
    )
    manager.initialize.assert_called_once_with()
    server.add_insecure_port.assert_called_once_with("[::]:19002")
    server.start.assert_called_once_with()
    assert resources.worker_endpoint == "trainer.test:19002"

    resources.close()


def test_resources_own_only_private_transport_resources():
    manager = MagicMock()
    manifest_service = MagicMock()
    server = MagicMock()
    resources = _TrainerResources(
        manager=manager,
        manifest_service=manifest_service,
        server=server,
        worker_endpoint="trainer.test:19000",
    )

    assert resources.manager is manager
    assert resources.manifest_service is manifest_service
    assert resources.worker_endpoint == "trainer.test:19000"
    resources.close()
    resources.close()

    server.stop.assert_called_once_with(grace=None)
    server.stop.return_value.wait.assert_called_once_with()
    manager.shutdown.assert_called_once_with()


def test_resources_are_not_part_of_public_api():
    assert not hasattr(modelexpress_rl, "TrainerResources")