From 633616d3f26702cc49ef41af66f3bf91d1a79c10 Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Thu, 11 Jun 2020 15:29:50 +0300
Subject: [PATCH] Update cleanup service example

---
 examples/services/cleanup_service.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/examples/services/cleanup_service.py b/examples/services/cleanup_service.py
index c18862a7..84806516 100644
--- a/examples/services/cleanup_service.py
+++ b/examples/services/cleanup_service.py
@@ -20,7 +20,7 @@ from glob import glob
 from shutil import rmtree
 from time import sleep, time
 
-from trains_agent import APIClient
+from trains.backend_api.session.client import APIClient
 
 from trains import Task
 
@@ -49,6 +49,10 @@ args = task.connect(args)
 
 
 # if we are running as a service, just enqueue ourselves into the services queue and let it run the optimization
+if args["run_as_service"] and task.running_locally():
+    verify = input('Stop local execution and execute remotely [y]/n ?').strip().lower()
+    args["run_as_service"] = not verify or verify.startswith('y')
+
 if args["run_as_service"]:
     # if this code is executed by `trains-agent` the function call does nothing.
     # if executed locally, the local process will be terminated, and a remote copy will be executed instead