Use Telemetry Configurations

Author: AryanAhadinia

src/swiss_ai_model_launch/assets/template.jinja

@@ -23,9 +23,11 @@ USE_OCF={{ "false" if disable_ocf else "true" }}
 OCF_SERVICE_NAME="llm"
 OCF_SERVICE_PORT=8080
 
-curl -sf -X POST "https://sml-dev.swissai.svc.cscs.ch/launches" \
+{% if telemetry_endpoint %}
+curl -sf -X POST "{{ telemetry_endpoint }}" \
     -H "Content-Type: application/json" \
     -d '{"user": "aryan", "job_id": "'"${SLURM_JOB_ID}"'", "slurm_nodes": {{ nodes }}, "slurm_job_name": "{{ job_name }}", "slurm_partition": "{{ partition }}", "slurm_time": "{{ time }}", "slurm_account": "{{ account }}", "slurm_environment": "{{ environment }}", "interactive": false, "serving_framework": "{{ framework }}", "framework_args": "{{ framework_args }}", "pre_launch_cmds": "{{ pre_launch_cmds }}", "model_name": "string", "model_path": "string", "tp_size": 0, "host": "string", "port": 0, "workers": {{ workers }}, "nodes_per_worker": {{ nodes_per_worker }}, "worker_port": {{ worker_port }}, "use_router": {{ "true" if use_router else "false" }}, "router_environment": "{{ environment }}", "router_port": 30000, "router_args": "{{ router_args }}", "ocf_enabled": {{ "false" if disable_ocf else "true" }}, "ocf_bootstrap_addr": "string", "ocf_service_name": "llm", "ocf_service_port": 8080, "extra": {"additionalProp1": {}}, "raw_command": "string"}' || true
+{% endif %}
 
 {% raw %}
 OCF_BOOTSTRAP_ADDR=$(curl -sf http://148.187.108.172:8092/v1/dnt/bootstraps \

src/swiss_ai_model_launch/cli/main.py

@@ -38,6 +38,7 @@ def _get_firecrest_client_from_init_config(config: InitConfig) -> f7t.v2.AsyncFi
 
 async def _get_firecrest_launcher_with_client(
     client: f7t.v2.AsyncFirecrest,
+    telemetry_endpoint: str | None = None,
 ) -> FirecRESTLauncher:
     async def _get_systems() -> dict[str, tuple[str, str]]:
         return {
@@ -82,6 +83,7 @@ async def _get_partitions(
         username=user_info["user"]["name"],
         account=user_info["group"]["name"],
         partition=launcher_config.get_non_none_value("firecrest_partition"),
+        telemetry_endpoint=telemetry_endpoint,
     )
 
 
@@ -225,7 +227,15 @@ async def _main() -> None:
     launcher_type = config.get_non_none_value("launcher")
     if launcher_type == "firecrest":
         firecrest_client = _get_firecrest_client_from_init_config(config)
-        launcher = await _get_firecrest_launcher_with_client(firecrest_client)
+        telemetry_endpoint = (
+            config.get_value("telemetry_endpoint")
+            if config.get_non_none_value("telemetry") == "default"
+            else None
+        )
+        launcher = await _get_firecrest_launcher_with_client(
+            firecrest_client,
+            telemetry_endpoint=telemetry_endpoint,
+        )
     else:
         raise NotImplementedError(f"Launcher {launcher_type} is not supported yet.")
 

src/swiss_ai_model_launch/launchers/firecrest_launcher.py

@@ -36,13 +36,15 @@ def __init__(
         username: str,
         account: str,
         partition: str,
+        telemetry_endpoint: str | None = None,
     ):
         super().__init__()
         self.client = client
         self.system_name = system_name
         self.username = username
         self.account = account
         self.partition = partition
+        self.telemetry_endpoint = telemetry_endpoint
 
     def _get_user_dir(self) -> str:
         return f"/users/{self.username}"
@@ -77,6 +79,7 @@ def _get_launch_args_from_request(
                 "--host 0.0.0.0 "
                 "--port 8080 "
             ),
+            telemetry_endpoint=self.telemetry_endpoint,
         )
 
     def _get_local_env_file_path(self, launch_request: LaunchRequest) -> str:

src/swiss_ai_model_launch/launchers/launch_args.py

@@ -21,6 +21,7 @@ class LaunchArgs(BaseModel):
     use_router: bool = False
     router_args: str = ""
     disable_ocf: bool = False
+    telemetry_endpoint: str | None = None
 
     @model_validator(mode="after")
     def set_defaults(self) -> "LaunchArgs":