62 lines
1.5 KiB
Python
62 lines
1.5 KiB
Python
import os
|
|
from multiprocessing import cpu_count, Pool
|
|
|
|
from django.core.management.base import BaseCommand
|
|
from django.db import connection
|
|
|
|
from .link_ranchers_parallel_worker import process_city
|
|
|
|
os.environ["DJANGO_ALLOW_ASYNC_UNSAFE"] = "true"
|
|
|
|
|
|
def run_worker(chunk):
|
|
import os
|
|
import django
|
|
|
|
os.environ.setdefault(
|
|
"DJANGO_SETTINGS_MODULE",
|
|
"Rasaddam_Backend.settings"
|
|
)
|
|
django.setup()
|
|
|
|
|
|
class Command(BaseCommand):
|
|
help = "Parallel link ranchers to cooperative by city"
|
|
|
|
def add_arguments(self, parser):
|
|
parser.add_argument(
|
|
'--worker',
|
|
type=int,
|
|
default=cpu_count() // 2
|
|
)
|
|
|
|
def handle(self, *args, **options):
|
|
from apps.herd.models import Rancher
|
|
workers = options['worker']
|
|
|
|
from multiprocessing import get_context
|
|
|
|
chunks = self.prepare_chunks()
|
|
|
|
ctx = get_context("spawn")
|
|
with ctx.Pool(processes=options["workers"]) as pool:
|
|
pool.map(run_worker, chunks)
|
|
|
|
city_ids = (
|
|
Rancher.objects.filter(city__isnull=False)
|
|
.values_list('city_id', flat=True)
|
|
.distinct()
|
|
)
|
|
|
|
self.stdout.write(
|
|
f"Starting parallel sync for {len(city_ids)} cities "
|
|
f"with {workers} workers"
|
|
)
|
|
|
|
connection.close()
|
|
|
|
with Pool(processes=workers) as pool:
|
|
pool.map(process_city, city_ids)
|
|
|
|
self.stdout.write(self.style.SUCCESS("DONE ✅"))
|