Merge from TIL-24/til-24-base@9352746

FINALS_README.md

@@ -6,7 +6,7 @@
    1. [Additional Directories](#additional-directories)
    2. [Simulator](#simulator)
 2. [Instructions](#instructions)
-   1. [Local Developement Environment Setup](#local-developement-environment-setup)
+   1. [Local Development Environment Setup](#local-development-environment-setup)
    2. [Artifacts Submission](#artifacts-submission)
    3. [Docker + Artifact Registry Recap](#docker-and-artifact-registry-recap)
    4. [Docker Compose](#docker-compose)
@@ -44,7 +44,7 @@ The simulator is written in JavaScript and uses `3js` to render the 3D environme
 
 ## Instructions
 
-### Local Developement Environment Setup
+### Local Development Environment Setup
 
 From here on, you will need to have a local development environment set up for your own testing. This local development environment should have:
 
@@ -105,7 +105,7 @@ When your model is pushed successfully, you should be able to see it under your
 
 ### Docker Compose
 
-To run all the containers, we will be using [Docker Compose](https://docs.docker.com/compose/), a tool that lets you easily configure and run multi-container applications.
+To run all the containers, we will be using [Docker Compose](https://docs.docker.com/compose/), a tool that lets you easily configure and run multi-container applications. Note that you will likely have to install it if it is not already installed (e.g. it may not be installed on the GCP instance); see [the Docker Compose plugin installation instructions](https://docs.docker.com/compose/install/#scenario-two-install-the-compose-plugin).
 
 ```bash
 # start all the services
@@ -181,3 +181,32 @@ pip install robomaster
 ```
 
 This will create a new Python 3.8 environment `robo` and configure it to use `x86` packages using Rosetta instead of native ones for `arm64`.
+
+### Loading Models Offline
+
+Some HuggingFace models, such as the DETR model referenced in the training curriculum, encounter some issues when trying to load from a finetuned model offline, as they rely on some backbone weights from HuggingFace Hub not stored in the model when downloaded using the `save_pretrained()` function.
+
+One possible solution to this problem is to run a script to load the models once when building the Dockerfile. This will cache whatever backbone weights are necessary, so subsequent calls of the model will just use the cached version instead of trying to access the HuggingFace Hub (which won't be possible without internet access).
+
+Here's an example of what such a script, `save_cache.py`, would look like:
+
+```python
+from transformers import (
+    AutoImageProcessor,
+    AutoModelForObjectDetection
+)
+
+detr_checkpoint = "models/detr_model.pth"
+
+detr_model = AutoModelForObjectDetection.from_pretrained(detr_checkpoint)
+detr_processor = AutoImageProcessor.from_pretrained(detr_checkpoint)
+```
+
+Then in your Dockerfile, add these lines after `pip install -r requirements.txt` and copying over necessary files (especially the model file).
+
+```dockerfile
+COPY save_cache.py save_cache.py
+RUN python3 save_cache.py
+```
+
+This will copy and run `save_cache.py`, which will save the model backbone weights to the `.cache` directory on the image, allowing it to be loaded offline by HuggingFace Transformers later on.

main/requirements.txt

@@ -1,3 +1,3 @@
 websockets
-requests
-asyncio
+asyncio
+httpx

main/src/auto_manager.py

@@ -1,41 +1,40 @@
 from random import randint
 from typing import Dict, List
 from finals_manager import FinalsManager
-import requests
 
 
 class AutoManager(FinalsManager):
     def __init__(self, local_ip: str):
         super().__init__()
         self.local_ip = local_ip
 
-    def run_asr(self, audio_bytes: bytes) -> str:
+    async def run_asr(self, audio_bytes: bytes) -> str:
         print("Running ASR")
         return "asr"
 
-    def run_nlp(self, transcript: str) -> Dict[str, str]:
+    async def run_nlp(self, transcript: str) -> Dict[str, str]:
         print("Running NLP")
         return {
             "target": "airplane",
             "heading": f"{randint(1,360):03}",
             "tool": "surface-to-air missiles",
         }
 
-    def run_vlm(self, image_bytes: bytes, caption: str) -> List[int]:
+    async def run_vlm(self, image_bytes: bytes, caption: str) -> List[int]:
         print("Running VLM")
         return [0, 0, 0, 0]
 
-    def send_heading(self, heading: str) -> bytes:
+    async def send_heading(self, heading: str) -> bytes:
         assert heading.isdigit(), "The heading string contains non-digit characters"
         print(f"Sending cannon heading {heading}")
-        results = requests.post(
+        results = await self.async_post(
             f"http://{self.local_ip}:5003/send_heading", json={"heading": heading}
         )
         # snapshot of image
         return results.content
 
-    def reset_cannon(self):
+    async def reset_cannon(self):
         print("Resetting cannon to original position")
-        results = requests.post(f"http://{self.local_ip}:5003/reset_cannon")
+        results = await self.async_post(f"http://{self.local_ip}:5003/reset_cannon")
         print(results.text)
         return results.json()

main/src/finals_manager.py

@@ -1,33 +1,42 @@
 import json
 import websockets
 from abc import ABC, abstractmethod
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
+import httpx
 
 
 class FinalsManager(ABC):
     def __init__(self):
         print("initializing participant finals server manager")
+        self.client = httpx.AsyncClient()
 
-    def send_result(
+    async def exit(self):
+        await self.client.aclose()
+
+    async def async_post(self, endpoint: str, json: Optional[Dict] = None):
+        return await self.client.post(endpoint, json=json)
+
+    async def send_result(
         self, websocket: websockets.WebSocketClientProtocol, data: Dict[str, Any]
     ):
-        return websocket.send(json.dumps(data))
+        return await websocket.send(json.dumps(data))
 
     @abstractmethod
-    def run_asr(self, audio_bytes: bytes) -> str:
+    async def run_asr(self, audio_bytes: bytes) -> str:
         raise NotImplemented
 
     @abstractmethod
-    def run_nlp(self, transcript: str) -> Dict[str, str]:
+    async def run_nlp(self, transcript: str) -> Dict[str, str]:
         raise NotImplemented
 
     @abstractmethod
-    def send_heading(self, heading: str) -> bytes:
+    async def send_heading(self, heading: str) -> bytes:
         raise NotImplemented
 
     @abstractmethod
-    def reset_cannon(self) -> None:
+    async def reset_cannon(self) -> None:
         raise NotImplemented
 
-    def run_vlm(self, image_bytes: bytes, caption: str) -> List[int]:
+    @abstractmethod
+    async def run_vlm(self, image_bytes: bytes, caption: str) -> List[int]:
         raise NotImplemented

main/src/mock_manager.py

@@ -1,4 +1,4 @@
-from time import sleep
+from asyncio import sleep
 from typing import Dict, List
 from finals_manager import FinalsManager
 
@@ -7,27 +7,27 @@ class MockManager(FinalsManager):
     def __init__(self):
         super().__init__()
 
-    def run_asr(self, audio_bytes: bytes) -> str:
+    async def run_asr(self, audio_bytes: bytes) -> str:
         print("Running ASR")
         return "asr"
 
-    def run_nlp(self, transcript: str) -> Dict[str, str]:
+    async def run_nlp(self, transcript: str) -> Dict[str, str]:
         print("Running NLP")
         return {
             "target": "airplane",
             "heading": "180",
             "tool": "surface-to-air missiles",
         }
 
-    def run_vlm(self, image_bytes: bytes, caption: str) -> List[int]:
+    async def run_vlm(self, image_bytes: bytes, caption: str) -> List[int]:
         print("Running VLM")
         return [0, 0, 0, 0]
 
-    def send_heading(self, heading: str) -> bytes:
+    async def send_heading(self, heading: str) -> bytes:
         print(f"Sending cannon heading {heading}")
-        sleep(1)
+        await sleep(1)
         return bytes()
 
-    def reset_cannon(self) -> None:
+    async def reset_cannon(self) -> None:
         print("Resetting cannon to original position")
         return {}

main/src/models_manager.py

@@ -1,48 +1,47 @@
 from base64 import b64encode
 from typing import Dict, List
 from finals_manager import FinalsManager
-import requests
 
 
 class ModelsManager(FinalsManager):
     def __init__(self, local_ip: str):
         super().__init__()
         self.local_ip = local_ip
 
-    def run_asr(self, audio_bytes: bytes) -> str:
+    async def run_asr(self, audio_bytes: bytes) -> str:
         print("Running ASR")
         audio_str = b64encode(audio_bytes).decode("ascii")
-        results = requests.post(
+        results = await self.async_post(
             f"http://{self.local_ip}:5001/stt", json={"instances": [{"b64": audio_str}]}
         )
         return results.json()["predictions"][0]
 
-    def run_nlp(self, transcript: str) -> Dict[str, str]:
+    async def run_nlp(self, transcript: str) -> Dict[str, str]:
         print("Running NLP")
-        results = requests.post(
+        results = await self.async_post(
             f"http://{self.local_ip}:5002/extract",
             json={"instances": [{"transcript": transcript}]},
         )
         return results.json()["predictions"][0]
 
-    def send_heading(self, heading: str) -> bytes:
+    async def send_heading(self, heading: str) -> bytes:
         assert heading.isdigit(), "The heading string contains non-digit characters"
         print(f"Sending cannon heading {heading}")
-        results = requests.post(
+        results = await self.async_post(
             f"http://{self.local_ip}:5003/send_heading", json={"heading": heading}
         )
         # snapshot of image
         return results.content
 
-    def reset_cannon(self):
+    async def reset_cannon(self):
         print("Resetting cannon to original position")
-        results = requests.post(f"http://{self.local_ip}:5003/reset_cannon")
+        results = await self.async_post(f"http://{self.local_ip}:5003/reset_cannon")
         return results.json()
 
-    def run_vlm(self, image_bytes: bytes, caption: str) -> List[int]:
+    async def run_vlm(self, image_bytes: bytes, caption: str) -> List[int]:
         print("Running VLM")
         image_str = b64encode(image_bytes).decode("ascii")
-        results = requests.post(
+        results = await self.async_post(
             f"http://{self.local_ip}:5004/identify",
             json={"instances": [{"b64": image_str, "caption": caption}]},
         )

main/src/participant_server.py

@@ -29,7 +29,7 @@ async def server():
         print(f"connecting to competition server {SERVER_IP} at port {SERVER_PORT}")
         try:
             while True:
-                # should be receiving either audio bytes for asr, or "done!" message
+                # should be receiving either audio bytes for asr, or done/healthcheck json
                 socket_input = await websocket.recv()
                 if type(socket_input) is str:
                     # handle either done or healthcheck
@@ -42,28 +42,28 @@ async def server():
                         continue
                 print(f"run {index}")
                 # ASR
-                transcript = manager.run_asr(socket_input)
+                transcript = await manager.run_asr(socket_input)
                 print(transcript)
                 # NLP
-                qa_ans = manager.run_nlp(transcript)
+                qa_ans = await manager.run_nlp(transcript)
                 print(qa_ans)
                 query = qa_ans["target"]
                 # autonomy
                 try:
-                    image = manager.send_heading(qa_ans["heading"])
+                    image = await manager.send_heading(qa_ans["heading"])
                 except AssertionError as e:
                     # if heading is wrong, get image of scene at default heading 000
                     print(e)
-                    image = manager.send_heading("000")
+                    image = await manager.send_heading("000")
                 # VLM
-                vlm_results = manager.run_vlm(image, query)
+                vlm_results = await manager.run_vlm(image, query)
                 print(vlm_results)
                 # submit results and reset
                 await manager.send_result(
                     websocket,
                     {"asr": transcript, "nlp": qa_ans, "vlm": vlm_results},
                 )
-                manager.reset_cannon()
+                await manager.reset_cannon()
                 print(f"done run {index}")
                 index += 1
         except websockets.ConnectionClosed: