Add CLI commands to start/stop web server. Also improve doc

Signed-off-by: Travis F. Collins <[email protected]>

bench/cli/iiotools.py

@@ -1,3 +1,7 @@
+import subprocess
+import os
+import time
+
 import click
 import requests
 from bench.keysight.dwta.data_capture import capture_iq_datafile
@@ -25,7 +29,7 @@ def cli(ctx, uri):
 @click.option(
     "--amplitude", "-a", help="Set the amplitude of the DDS in 0->1", required=True
 )
-@click.option("--device", "-d", help="Device driver to use")
+@click.option("--device", "-d", help="IIO device driver name to use")
 @click.option("--channel", "-c", help="Set the channel of the DDS", required=True)
 @click.option("--complex", "-x", is_flag=True, help="Use complex mode")
 @click.pass_context
@@ -77,12 +81,29 @@ def set_dds(ctx, frequency, amplitude, device, channel, complex):
 @cli.command()
 @click.option("--filename", "-f", help="Name of file to write data to", required=True)
 @click.option("--device", "-d", help="Name of device to configure", required=True)
-@click.option("--channel", "-c", help="Channel to capture data from", required=True)
+@click.option(
+    "--channel",
+    "-c",
+    help="Channel index to capture data from. Starts from 0",
+    required=True,
+)
 @click.option("--samples", "-s", help="Number of samples to capture", required=True)
-@click.argument("props", nargs=-1)
+@click.argument(
+    "props",
+    nargs=-1,
+    required=False,
+)
 @click.pass_context
-def capture_Data(ctx, filename, device, channel, samples, props):
+def capture_data(ctx, filename, device, channel, samples, props):
+    """Capture IQ data to a file in DWTA format
+
+    PROPS is a list of property=value pairs to set device properties. These are
+    the properties available in the pyadi-iio class interface for the device.
+
+    Example usage with ADALM-PLUTO:
 
+        pybenchiio -u ip:analog.local capture_data -f data.csv -d Pluto -c 0 -s 1024 sample_rate=1000000
+    """
     # Checks
     samples = int(samples)
     channel = int(channel)
@@ -118,11 +139,21 @@ def capture_Data(ctx, filename, device, channel, samples, props):
     default="localhost",
 )
 @click.option(
-    "--server-port", "-p", help="Port of the server", required=False, default=8000
+    "--server-port", "-p", help="Port of the server", required=False, default=12345
 )
 @click.argument("props", nargs=-1)
 @click.pass_context
 def transmit_data(ctx, filename, device, channel, server_ip, server_port, props):
+    """Transmit IQ data file to device through backend server
+
+    File must be in DWTA format, where the first two lines are sample rate and
+    center frequency, and the rest of the lines are IQ data in the format
+    I, Q per line.
+
+    Example usage with ADALM-PLUTO:
+
+        pybenchiio -u ip:analog.local transmit_data -f data.csv -d Pluto -c 0 sample_rate=1000000
+    """
 
     # Checks
     channel = int(channel)
@@ -173,10 +204,94 @@ def transmit_data(ctx, filename, device, channel, server_ip, server_port, props)
     default="localhost",
 )
 @click.option(
-    "--server-port", "-p", help="Port of the server", required=False, default=8000
+    "--server-port", "-p", help="Port of the server", required=False, default=12345
 )
 def transmit_data_clear(server_ip, server_port):
+    """Clear the transmit buffer on the server"""
 
     url = f"http://{server_ip}:{server_port}/clearbuffer"
     r = requests.post(url)
     assert r.status_code == 200, f"Failed to clear buffer: {r.json()}"
+
+
+@cli.command()
+@click.option(
+    "--host",
+    "-h",
+    help="Host to start the server on",
+    required=False,
+    default="localhost",
+)
+@click.option(
+    "--port", "-p", help="Port to start the server on", required=False, default=12345
+)
+def start_server(host, port):
+
+    # Start the server as a subprocess
+    import subprocess
+    import os
+
+    loc = os.path.dirname(os.path.realpath(__file__))
+    web_app = os.path.join(loc, "..", "web", "app.py")
+    web_app = os.path.abspath(web_app)
+
+    # Get location of python executable
+    python = "python"
+    if "CONDA_PREFIX" in os.environ:
+        python = os.path.join(os.environ["CONDA_PREFIX"], "bin", "python")
+    elif "VIRTUAL_ENV" in os.environ:
+        python = os.path.join(os.environ["VIRTUAL_ENV"], "bin", "python")
+    elif "PYTHON" in os.environ:
+        python = os.environ["PYTHON"]
+
+    command = [
+        python,
+        "-m",
+        "fastapi",
+        "run",
+        web_app,
+        "--host",
+        str(host),
+        "--port",
+        str(port),
+    ]
+    # Start process and verify running after 5 seconds
+    p = subprocess.Popen(command)
+    time.sleep(5)
+    assert p.poll() is None, "Failed to start server"
+    print(f"Server started on {host}:{port} with PID {p.pid}")
+
+
+# Stop server
+@cli.command()
+@click.option(
+    "--server-port", "-p", help="Port of the server", required=False, default=12345
+)
+def stop_server(server_port):
+    """Stop the server"""
+    # Find the process using the port and kill it
+    if os.name == "nt":
+        # Query for the process ID
+        pid = subprocess.run(
+            f"netstat -aon | findstr {server_port}", shell=True, stdout=subprocess.PIPE
+        )
+        if pid.returncode != 0:
+            print("Server not running")
+            return
+        pid = pid.stdout.decode().split(" ")[-1]
+        # Kill the process
+        subprocess.run(f"taskkill /F /PID {pid}", shell=True)
+    else:
+        # Query for the process ID
+        pid = subprocess.run(
+            f"lsof -t -i:{server_port}", shell=True, stdout=subprocess.PIPE
+        )
+        if pid.returncode != 0:
+            print("Server not running")
+            return
+        pid = pid.stdout.decode().strip()
+        print(f"Server found on port {server_port} with PID {pid}")
+        # Kill the process
+        subprocess.run(f"kill -9 {pid}", shell=True)
+
+    print(f"Server on port {server_port} stopped")

bench/web/app.py

@@ -1,4 +1,5 @@
 """Web backend for the bench"""
+
 import os
 from typing import List, Union
 import threading
@@ -16,21 +17,26 @@
 # Import supported_devices from the data_capture module
 from bench.keysight.dwta.data_capture import supported_devices
 
+
 class SharedState:
     def __init__(self):
         self.buffer = None
         self.device = None
 
+
 state = SharedState()
 
+
 def read_state(prop):
     global state
     return getattr(state, prop)
 
+
 def write_state(prop, value):
     global state
     setattr(state, prop, value)
 
+
 class BufferWrite(BaseModel):
     uri: str
     device: str
@@ -47,29 +53,24 @@ class BufferWrite(BaseModel):
 file_location = os.path.dirname(os.path.realpath(__file__))
 templates = Jinja2Templates(directory=os.path.join(file_location, "templates"))
 
-@app.get("/")
-def read_root():
-    return {"Hello": "World"}
 
-
-@app.get("/items/{item_id}")
-def read_item(item_id: int, q: Union[str, None] = None):
-    return {"item_id": item_id, "q": q}
-
-@app.post('/clearbuffer')
+@app.post("/clearbuffer")
 async def clearbuffer():
     write_state("buffer", None)
     write_state("device", None)
     return {"status": "ok"}
 
+
 @app.post("/writebuffer")
 async def writebuffer(bufferwrite: BufferWrite):
     # print(bufferwrite)
     # print(bufferwrite.channels)
 
     # Checks
     if bufferwrite.device not in supported_devices:
-        return {"status": f"Device not supported: {bufferwrite.device}. Supported devices: {supported_devices}"}
+        return {
+            "status": f"Device not supported: {bufferwrite.device}. Supported devices: {supported_devices}"
+        }
 
     # Create device
     device = getattr(adi, bufferwrite.device)(bufferwrite.uri)
@@ -105,13 +106,15 @@ async def writebuffer(bufferwrite: BufferWrite):
     # Save device state
     write_state("device", device)
 
-
     return {"status": "ok"}
 
+
 @app.get("/help/{id}", response_class=HTMLResponse)
 async def help(request: Request, id: str):
-    return templates.TemplateResponse(request=request, name="help.html", context={"id": id})
+    return templates.TemplateResponse(
+        request=request, name="help.html", context={"id": id}
+    )
 
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    uvicorn.run(app, host="0.0.0.0", port=8000)

pyproject.toml

@@ -50,7 +50,7 @@ pybenchiio = "bench.cli.iiotools:cli"
 pybench = "bench.cli.visatools:cli"
 
 [project.optional-dependencies]
-cli = ["click", "pyadi-iio>=0.0.17"]
+cli = ["click", "pyadi-iio>=0.0.17", "requests"]
 web = ["fastapi[standard]", "uvicorn", "pydantic", "jinja2", "starlette"]
 
 [project.entry-points.pytest11]