Skip to content

Commit ab80569

Browse files
raghav-nexthopraghav-nexthop
committed
Distro CLI: Implement DeviceUpdater update logic
- Implement full update workflow in DeviceUpdater.update() - Add artifact acquisition via build or download - Add update package creation with update_service.sh script - Add SSH-based transfer and execution on device - Add integration with distro_infra for device IP resolution - Add unit tests for update workflow - Add fboss_init.sh support for distro-base snapshot creation
1 parent 56cb9f6 commit ab80569

File tree

7 files changed

+361
-44
lines changed

7 files changed

+361
-44
lines changed

‎cmake/FbossImageDistroCliTests.cmake‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@ find_package(Python3 3.10 COMPONENTS Interpreter REQUIRED)
1717
message(STATUS "Using Python ${Python3_VERSION} (${Python3_EXECUTABLE}) for distro_cli tests")
1818

1919
include(FBPythonBinary)
20-
2120
file(GLOB DISTRO_CLI_TEST_SOURCES
2221
"fboss-image/distro_cli/tests/*_test.py"
2322
)
@@ -86,6 +85,19 @@ add_custom_command(
8685
COMMENT "Copying test data files for distro_cli_tests"
8786
)
8887

88+
# Copy scripts directory used in unit tests
89+
set(SCRIPTS_DEST_DIR "${CMAKE_CURRENT_BINARY_DIR}/distro_cli_tests/distro_cli/scripts")
90+
set(SCRIPTS_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/fboss-image/distro_cli/scripts")
91+
92+
add_custom_command(
93+
TARGET distro_cli_tests.GEN_PY_EXE
94+
POST_BUILD
95+
COMMAND ${CMAKE_COMMAND} -E copy_directory
96+
"${SCRIPTS_SOURCE_DIR}"
97+
"${SCRIPTS_DEST_DIR}"
98+
COMMENT "Copying scripts for distro_cli_tests"
99+
)
100+
89101
install_fb_python_executable(distro_cli_tests)
90102

91103
# Restore the original Python3_EXECUTABLE if it was set

‎fboss-image/distro_cli/cmds/device.py‎

Lines changed: 56 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,10 @@
1111
import logging
1212
import os
1313
import sys
14+
from pathlib import Path
1415

1516
from distro_cli.lib.cli import validate_path
17+
from distro_cli.lib.device_update import DeviceUpdateError, DeviceUpdater
1618
from distro_cli.lib.distro_infra import (
1719
DISTRO_INFRA_CONTAINER,
1820
GETIP_SCRIPT_CONTAINER_PATH,
@@ -21,6 +23,7 @@
2123
)
2224
from distro_cli.lib.docker import container
2325
from distro_cli.lib.exceptions import DistroInfraError
26+
from distro_cli.lib.manifest import ImageManifest
2427

2528
logger = logging.getLogger(__name__)
2629

@@ -57,16 +60,66 @@ def image_command(args):
5760

5861
def reprovision_command(args):
5962
"""Reprovision device"""
60-
logger.info(f"Reprovisioning device {args.mac}")
61-
logger.info("Device reprovision command (stub)")
63+
ip_address = get_device_ip(args.mac)
64+
65+
if not ip_address:
66+
logger.error("No IP address found for device")
67+
return
68+
69+
# devpart -> /dev/nvme0n1p3
70+
# dev -> /dev/nvme0n3
71+
# part -> 3
72+
cmd = r"""
73+
if [ ! -d /opt/fboss ]; then echo "Not an FBOSS device. Aborting"; exit 1; fi; \
74+
rm -rf /boot/efi/EFI/*;
75+
root_devpart=$(mount | awk '/\/ type/ { print $1 }');
76+
root_dev=$(mount | awk -F 'p' '/\/ type/ { print $1 }');
77+
root_part=$(mount | awk -F '[[:space:]p]' '/\/ type/ { print $2 }');
78+
dd if=/dev/zero of=${root_devpart} bs=1M count=50;
79+
(sleep 1; echo yes; sleep 1; echo ignore) | parted ---pretend-input-tty ${root_dev} rm ${root_part};
80+
reboot --force
81+
"""
82+
os.execvp(
83+
"ssh",
84+
[
85+
"ssh",
86+
"-o",
87+
"StrictHostKeyChecking=no",
88+
"-o",
89+
"UserKnownHostsFile=/dev/null",
90+
f"root@{ip_address}",
91+
cmd,
92+
],
93+
)
6294

6395

6496
def update_command(args):
6597
"""Update specific components on device"""
6698
logger.info(f"Updating device {args.mac}")
6799
logger.info(f"Manifest: {args.manifest}")
68100
logger.info(f"Components: {' '.join(args.components)}")
69-
logger.info("Device update command (stub)")
101+
102+
manifest = ImageManifest(Path(args.manifest))
103+
104+
# Get device IP once for all components
105+
device_ip = get_device_ip(args.mac)
106+
if not device_ip:
107+
logger.error("Cannot update: device IP not found")
108+
sys.exit(1)
109+
110+
for component in args.components:
111+
try:
112+
updater = DeviceUpdater(
113+
mac=args.mac,
114+
manifest=manifest,
115+
component=component,
116+
device_ip=device_ip,
117+
)
118+
updater.update()
119+
logger.info(f"Successfully updated {component}")
120+
except DeviceUpdateError as e:
121+
logger.error(f"Failed to update {component}: {e}")
122+
sys.exit(1)
70123

71124

72125
def get_device_ip(mac: str) -> str | None:

‎fboss-image/distro_cli/lib/device_update.py‎

Lines changed: 120 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,11 @@
88
"""Device update logic for updating FBOSS services on devices."""
99

1010
import logging
11+
import subprocess
12+
import uuid
1113
from pathlib import Path
1214

15+
from distro_cli.builder.image_builder import ImageBuilder
1316
from distro_cli.lib.exceptions import DistroInfraError
1417
from distro_cli.lib.manifest import ImageManifest
1518

@@ -26,6 +29,9 @@
2629
],
2730
}
2831

32+
# Path to the update script that runs on the device
33+
UPDATE_SCRIPT_PATH = Path(__file__).parent.parent / "scripts" / "update_service.sh"
34+
2935

3036
class DeviceUpdateError(DistroInfraError):
3137
"""Error during device update."""
@@ -37,10 +43,8 @@ class DeviceUpdater:
3743
Workflow:
3844
1. Validate component is supported for update
3945
2. Acquire artifacts (build OR download)
40-
3. Create update package (artifacts + service_update.sh) to scp to the device
41-
4. Get device IP
42-
5. SCP update package to device
43-
6. SSH: extract and run service_update.sh
46+
3. SCP artifact and update_service.sh to device
47+
4. SSH: run update_service.sh
4448
"""
4549

4650
def __init__(
@@ -102,39 +106,121 @@ def validate(self) -> None:
102106
def _acquire_artifacts(self) -> Path:
103107
"""Acquire component artifacts via build or download.
104108

109+
Uses ImageBuilder to handle both execute (build) and download modes.
110+
Dependencies are automatically built if needed.
111+
105112
Returns:
106113
Path to the component artifact (tarball)
107114

108115
Raises:
109116
DeviceUpdateError: If artifact acquisition fails
110117
"""
111-
raise NotImplementedError("Stub")
112-
113-
def _create_update_package(self, artifact_path: Path) -> Path:
114-
"""Create update package with artifacts and update_service.sh script.
118+
logger.info(f"Acquiring artifacts for {self.component}")
115119

116-
Args:
117-
artifact_path: Path to the component artifact tarball
120+
builder = ImageBuilder(self.manifest)
121+
builder.build_components([self.component])
118122

119-
Returns:
120-
Path to the created update package
123+
artifact_path = builder.component_artifacts.get(self.component)
124+
if not artifact_path:
125+
raise DeviceUpdateError(
126+
f"No artifact produced for component '{self.component}'"
127+
)
121128

122-
Raises:
123-
DeviceUpdateError: If package creation fails
124-
"""
125-
raise NotImplementedError("Stub")
129+
logger.info(f"Artifact acquired: {artifact_path}")
130+
return artifact_path
126131

127-
def _transfer_and_execute(self, package_path: Path, services: list[str]) -> None:
128-
"""Transfer update package to device and execute update_service.sh.
132+
def _transfer_and_execute(self, artifact_path: Path, services: list[str]) -> None:
133+
"""Transfer artifact and update script to device and execute.
129134

130135
Args:
131-
package_path: Path to the update package tarball
136+
artifact_path: Path to the component artifact tarball
132137
services: List of services to update
133138

134139
Raises:
135140
DeviceUpdateError: If transfer or execution fails
136141
"""
137-
raise NotImplementedError("Stub")
142+
if not self.device_ip:
143+
raise DeviceUpdateError("Device IP not set")
144+
145+
if not UPDATE_SCRIPT_PATH.exists():
146+
raise DeviceUpdateError(f"Update script not found: {UPDATE_SCRIPT_PATH}")
147+
148+
device_user = "root"
149+
remote_dir = f"/tmp/fboss-update-{uuid.uuid4().hex[:8]}"
150+
ssh_opts = [
151+
"-o",
152+
"StrictHostKeyChecking=no",
153+
"-o",
154+
"UserKnownHostsFile=/dev/null",
155+
]
156+
157+
logger.info(f"Transferring files to {self.device_ip}:{remote_dir}")
158+
159+
try:
160+
# Create remote directory
161+
result = subprocess.run(
162+
[
163+
"ssh",
164+
*ssh_opts,
165+
f"{device_user}@{self.device_ip}",
166+
f"mkdir -p {remote_dir}",
167+
],
168+
capture_output=True,
169+
check=False,
170+
)
171+
if result.returncode != 0:
172+
raise DeviceUpdateError(
173+
f"Failed to create remote directory: {result.stderr.decode()}"
174+
)
175+
176+
# SCP artifact and update script to device
177+
result = subprocess.run(
178+
[
179+
"scp",
180+
*ssh_opts,
181+
str(artifact_path),
182+
str(UPDATE_SCRIPT_PATH),
183+
f"{device_user}@{self.device_ip}:{remote_dir}/",
184+
],
185+
capture_output=True,
186+
check=False,
187+
)
188+
if result.returncode != 0:
189+
raise DeviceUpdateError(
190+
f"Failed to transfer files: {result.stderr.decode()}"
191+
)
192+
193+
services_arg = " ".join(services)
194+
remote_cmd = (
195+
f"cd {remote_dir} && "
196+
f"chmod +x update_service.sh && "
197+
f"sudo ./update_service.sh {self.component} {services_arg}"
198+
)
199+
200+
logger.info("Executing update on device...")
201+
result = subprocess.run(
202+
["ssh", *ssh_opts, f"{device_user}@{self.device_ip}", remote_cmd],
203+
capture_output=True,
204+
check=False,
205+
)
206+
if result.returncode != 0:
207+
raise DeviceUpdateError(
208+
f"Failed to execute update: {result.stderr.decode()}"
209+
)
210+
211+
logger.info(f"Update output:\n{result.stdout.decode()}")
212+
finally:
213+
# Cleanup remote directory if present
214+
subprocess.run(
215+
[
216+
"ssh",
217+
*ssh_opts,
218+
f"{device_user}@{self.device_ip}",
219+
f"rm -rf {remote_dir}",
220+
],
221+
capture_output=True,
222+
check=False,
223+
)
138224

139225
def update(self) -> bool:
140226
"""Execute the update workflow.
@@ -145,4 +231,17 @@ def update(self) -> bool:
145231
Raises:
146232
DeviceUpdateError: If update fails
147233
"""
148-
raise NotImplementedError("Stub")
234+
if not self.device_ip:
235+
raise DeviceUpdateError("Device IP not set")
236+
237+
self.validate()
238+
239+
services = self._get_services()
240+
logger.info(f"Updating {self.component} on device {self.mac}")
241+
logger.info(f"Services to restart: {', '.join(services)}")
242+
243+
artifact_path = self._acquire_artifacts()
244+
self._transfer_and_execute(artifact_path, services)
245+
246+
logger.info(f"Successfully updated {self.component} on device {self.mac}")
247+
return True

‎fboss-image/distro_cli/scripts/update_service.sh‎

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
#!/bin/bash
2+
# Copyright (c) 2004-present, Facebook, Inc.
3+
# All rights reserved.
4+
#
5+
# This source code is licensed under the BSD-style license found in the
6+
# LICENSE file in the root directory of this source tree.
7+
8+
# Update script that runs on the FBOSS device to update services.
9+
# Installs whatever artifact exists in the script's directory.
10+
# Usage: ./update_service.sh <component> <service1> [service2] ...
11+
12+
set -eou pipefail
13+
14+
if [ $# -lt 2 ]; then
15+
echo "Usage: $0 <component> <service1> [service2] ..." >&2
16+
exit 1
17+
fi
18+
19+
COMPONENT="$1"
20+
shift
21+
SERVICES="$*"
22+
23+
SCRIPT_DIR=$(dirname "$(readlink -f "$0")")
24+
TIMESTAMP=$(date +%s)
25+
BASE_SNAPSHOT="/distro-base"
26+
UPDATES_DIR="/updates"
27+
28+
echo "Updating component: ${COMPONENT}"
29+
echo "Services to restart: ${SERVICES}"
30+
31+
# Extract all artifacts from script's directory to staging
32+
mkdir -p "${UPDATES_DIR}"
33+
STAGING_DIR=$(mktemp -d -p "${UPDATES_DIR}")
34+
trap 'rm -rf "${STAGING_DIR}"' EXIT
35+
36+
FOUND_ARTIFACT=false
37+
for f in "${SCRIPT_DIR}"/*.tar.zst "${SCRIPT_DIR}"/*.tar; do
38+
FOUND_ARTIFACT=true
39+
echo "Extracting ${f}..."
40+
tar -xf "$f" -C "${STAGING_DIR}"
41+
done
42+
43+
if [ "${FOUND_ARTIFACT}" = false ]; then
44+
echo "Error: No artifact found in ${SCRIPT_DIR}" >&2
45+
exit 1
46+
fi
47+
48+
for svc in ${SERVICES}; do
49+
SNAPSHOT_PATH="${UPDATES_DIR}/${svc}-${TIMESTAMP}"
50+
echo "Creating snapshot for ${svc}: ${SNAPSHOT_PATH}"
51+
52+
btrfs subvolume snapshot "${BASE_SNAPSHOT}" "${SNAPSHOT_PATH}"
53+
cp -a "${STAGING_DIR}"/* "${SNAPSHOT_PATH}/opt/fboss/"
54+
55+
mkdir -p "/etc/systemd/system/${svc}.service.d/"
56+
cat >"/etc/systemd/system/${svc}.service.d/root-override.conf" <<EOF
57+
[Service]
58+
RootDirectory=${SNAPSHOT_PATH}
59+
EOF
60+
done
61+
62+
echo "Reloading systemd and restarting services..."
63+
systemctl daemon-reload
64+
for svc in ${SERVICES}; do
65+
echo " Restarting ${svc}..."
66+
systemctl restart "${svc}"
67+
done
68+
69+
# Delete old subvolumes after restart succeeds
70+
for svc in ${SERVICES}; do
71+
SNAPSHOT_PATH="${UPDATES_DIR}/${svc}-${TIMESTAMP}"
72+
for old in "${UPDATES_DIR}/${svc}"-*; do
73+
if [ -d "$old" ] && [ "$old" != "${SNAPSHOT_PATH}" ]; then
74+
echo " Deleting old snapshot: ${old}"
75+
btrfs subvolume delete "$old" 2>/dev/null || true
76+
fi
77+
done
78+
done
79+
80+
echo "Update complete for component: ${COMPONENT}"

‎fboss-image/distro_cli/tests/data/artifacts/platform_stack.tar.zst‎

115 Bytes
Binary file not shown.

0 commit comments

Comments
 (0)